Compare commits

...

6 commits

Author SHA1 Message Date
Yiorgis Gozadinos
1ecf77cad5
Hooks own resources through a lifespan 2026-08-24 18:02:23 +03:00
Yiorgis Gozadinos
1cb5e296b0
Hooks build from the resolved client config
HaikuRAG resolves `config=None` to the global config for everything
else; the hook list read the raw argument, so a client constructed
without an explicit config never loaded its configured hooks.

Tests follow the config and context-expansion APIs: `get_config()`
replaces the removed `Config` instance, and the expansion test does
its own `resolve_refs_grouped` / `get_items_in_ranges` fetch now that
`expand_with_items` takes the items rather than fetching them.
2026-08-24 18:02:23 +03:00
Yiorgis Gozadinos
e10762854c
Post-commit hooks are best-effort observers 2026-08-24 18:02:23 +03:00
Yiorgis Gozadinos
8083c57246
Hook events carry batches and search parameters 2026-08-24 18:02:22 +03:00
Yiorgis Gozadinos
b4a3c678a0
Add SearchResult.annotations for after_search hooks 2026-08-24 18:02:09 +03:00
Yiorgis Gozadinos
330468f2ca
Add client lifecycle hooks (after_ingest, after_delete, before_search, after_search) 2026-08-24 18:02:09 +03:00
15 changed files with 1263 additions and 31 deletions

View file

@ -1,6 +1,12 @@
# 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. Post-commit observer hooks are best-effort; failures are logged and do not change the completed operation's result.
- `SearchResult.annotations` carries notes attached by `after_search` hooks, preserved through context expansion and rendered in agent-facing output.
- `Hook.lifespan`, an async context manager around the client's lifetime for hooks that own resources. Entered in configured order once the store is open, exited in reverse order while the store, embedder and reranker are still usable.
## [0.78.0] - 2026-08-24
### Added

View file

@ -31,6 +31,7 @@ Built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/
- **MCP server** — Expose as tools for AI assistants (Claude Desktop, etc.)
- **Visual grounding** — View chunks highlighted on original page images
- **Production ingester** — Long-lived `haiku-ingester` service with persistent SQLite queue, async worker pool with retries and a dead-letter queue, FS / HTTP / S3 / WebDAV source adapters, FastAPI control plane, and a browser dashboard for operators. See [docs/ingester.md](docs/ingester.md).
- **Hooks** — Plugin packages can observe document writes and transform searches (query expansion, result annotation) via the `haiku.rag.hooks` entry-point group. See [docs/hooks.md](docs/hooks.md).
- **Tags** — Name database states with `haiku-rag tag` and roll back to them
- **Inspector** — TUI for browsing documents, chunks, and search results

View file

@ -71,6 +71,8 @@ qa:
# haiku.rag.yaml
environment: production
hooks: [] # Lifecycle hook plugin names, see the Hooks page under Develop
storage:
data_dir: "" # Empty = use default platform location
vacuum_retention_seconds: 86400

111
docs/hooks.md Normal file
View file

@ -0,0 +1,111 @@
# Hooks
Hooks let external packages observe document writes and transform searches without forking haiku.rag. Use them for query rewriting, result annotation, or maintaining state derived from the corpus (a synonym table, an entity index, corpus statistics).
A hook is a class registered under the `haiku.rag.hooks` entry-point group and activated by name in config. Hooks run everywhere the client runs: CLI, MCP server, skills, and your own code.
## Hook points
Subclass `haiku.rag.hooks.Hook` and override any subset:
| Method | Fires | Use for |
|--------|-------|---------|
| `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 |
| `lifespan(client)` | Around the whole client lifetime | Owning connections, clients, background tasks |
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, and reads a request whose `search_type` is the one retrieval actually used: `hybrid` where a text search left it unset, `vector` for an image query. 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.
## Registering a hook
```python
from haiku.rag.hooks import Hook
class AbbreviationHook(Hook):
async def before_search(self, client, request):
request.query = my_glossary.expand(request.query)
return request
async def after_search(self, client, request, results):
for result in results:
result.annotations = [
f"{term}: {definition}"
for term, definition in my_glossary.definitions_in(result.content)
]
return results
```
Register a zero-arg factory in your package's `pyproject.toml`:
```toml
[project.entry-points."haiku.rag.hooks"]
abbreviations = "my_package.hooks:AbbreviationHook"
```
Activate it in `haiku.rag.yaml`:
```yaml
hooks:
- abbreviations
```
An unknown name in `hooks:` raises `ValueError` when the client is constructed, so misconfiguration fails at startup. Entry points load lazily. Only the hooks the config references are imported.
## Owning resources
Factories are called during client construction, before the database is open, so they must not acquire resources. Acquire them in `lifespan` instead, an async context manager around the client's lifetime:
```python
from contextlib import asynccontextmanager
import httpx
from haiku.rag.hooks import Hook
class GlossaryHook(Hook):
@asynccontextmanager
async def lifespan(self, client):
async with httpx.AsyncClient() as http:
self.http = http
yield
```
Lifespans are entered in the order listed in config, once the store is open, and exited in reverse order while the store, embedder and reranker are all still usable.
Failing on entry fails `async with HaikuRAG(...)` and unwinds the lifespans already started: an activated hook that cannot start is a startup failure, not something to run degraded. Failing on exit is logged and swallowed, so one hook's teardown cannot strand another's. A hook is told which exception is being unwound, whether it came from the client's caller or from a later hook failing to start, but cannot suppress it.
## Background work
A hook that runs a background task owns stopping it. Cancel it on the way out, before anything awaits it:
```python
@asynccontextmanager
async def lifespan(self, client):
async with asyncio.TaskGroup() as tg:
task = tg.create_task(self.refresh_periodically(client))
try:
yield
finally:
task.cancel()
```
The `task.cancel()` is not optional. A `TaskGroup` that exits cleanly waits for its children instead of cancelling them, so an endless loop parked in one never returns and client shutdown hangs.
## Result annotations
`after_search` hooks can attach free-text notes on `SearchResult.annotations`. Annotations survive context expansion (merged results union the notes of their constituents, deduplicated) and render as `Note: text` lines in the agent-facing output used by the QA skills. MCP responses carry the field as part of the `SearchResult` model. This keeps the context cost proportional to what was retrieved instead of the size of your vocabulary.
## Semantics
- **Post-commit hooks are best-effort observers.** By the time `after_ingest` or `after_delete` runs, the operation has committed. A hook failure is logged, subsequent hooks still run, and the operation still returns success (the ingester proceeds through its normal success path). Correctness-critical derived state therefore needs its own retry or reconciliation, such as the backfill loop below. `before_search` and `after_search` failures propagate: nothing has committed and failing the search is visible to the caller.
- **Post-commit hooks are not a supported transformation point.** Mutating event models does not alter the committed record; explicit client writes are separate operations and are not atomic with the original write.
- **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`.
- **State lives in the database.** Hooks may create their own LanceDB tables via `client.store`. Prefix table names with `hook_` so they never collide with core tables or future migrations. State then travels with the database and its backups.

View file

@ -5,6 +5,7 @@ import logging
import mimetypes
import tempfile
from collections.abc import AsyncGenerator, Sequence
from contextlib import AsyncExitStack
from enum import Enum
from functools import cached_property
from pathlib import Path
@ -17,6 +18,13 @@ 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 (
DeleteEvent,
build_hooks,
enter_lifespans,
load_hooks,
notify,
)
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 +101,9 @@ class HaikuRAG:
self._vacuum_tasks: set[asyncio.Task] = set()
self._last_vacuum_at: float | None = None
self._vacuum_dirty = False
self._hooks = (
build_hooks(self._config.hooks, load_hooks()) if self._config.hooks else []
)
@property
def is_read_only(self) -> bool:
@ -114,33 +125,54 @@ class HaikuRAG:
return get_reranker(config=self._config)
async def __aenter__(self):
"""Async context manager entry — initializes store and repositories."""
self.store = Store(
self._db_path,
config=self._config,
skip_validation=self._skip_validation,
create=self._create,
read_only=self._read_only,
)
# If _initialize fails mid-way (e.g. migration check raises after
# connect), close the store so we don't leak the LanceDB connection —
# __aexit__ won't run because the `async with` never entered.
"""Async context manager entry — opens the store, the repositories and
the hook lifespans on one stack, so a failure anywhere along the way
unwinds whatever is already open. ``__aexit__`` won't run when the
`async with` never entered, so entry unwinds its own stack."""
stack = AsyncExitStack()
try:
self.store = Store(
self._db_path,
config=self._config,
skip_validation=self._skip_validation,
create=self._create,
read_only=self._read_only,
)
stack.callback(self.close)
await self.store._initialize()
except BaseException:
self.store.close()
self.document_repository = DocumentRepository(self.store)
self.chunk_repository = ChunkRepository(self.store)
self.document_item_repository = DocumentItemRepository(self.store)
stack.push_async_callback(self._close_models)
stack.push_async_callback(self._await_vacuum_tasks)
await enter_lifespans(stack, self._hooks, self)
except BaseException as exc:
# Forwarded, not aclose()'d: a hook that fails to start is an
# unwind like any other, and the lifespans already running are
# entitled to know what it was.
await stack.__aexit__(type(exc), exc, exc.__traceback__)
raise
self.document_repository = DocumentRepository(self.store)
self.chunk_repository = ChunkRepository(self.store)
self.document_item_repository = DocumentItemRepository(self.store)
self._stack = stack
return self
async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002
"""Async context manager exit."""
await self._await_vacuum_tasks()
# Best-effort: __aexit__ may run during exception unwinding, and a
# raising close must not mask the original exception. The reranker is
# a cached_property — close it only if it was materialized.
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Async context manager exit — unwinds hook lifespans in reverse
order, then the embedder and reranker, then the store.
The exception being unwound is forwarded to the lifespans so a hook can
tell a clean shutdown from a failing one. Suppressing it is not theirs
to decide, so the result is discarded.
"""
await self._stack.__aexit__(exc_type, exc_val, exc_tb)
return False
async def _close_models(self) -> None:
"""Close the embedder and the reranker.
Best-effort: teardown may run during exception unwinding, and a raising
close must not mask the original exception. The reranker is a
cached_property close it only if it was materialized.
"""
try:
await self.embedder.aclose()
reranker = self.__dict__.get("reranker")
@ -148,8 +180,6 @@ class HaikuRAG:
await reranker.aclose()
except Exception:
logger.debug("Closing embedder/reranker failed on teardown", exc_info=True)
self.close()
return False
async def _await_vacuum_tasks(self) -> None:
"""Drain background vacuum work and run a final collapse before teardown.
@ -398,7 +428,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:
@ -406,20 +436,23 @@ 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()
event = DeleteEvent(documents=docs_to_delete)
await notify(self._hooks, "after_delete", self, event)
return True
async def list_documents(

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, notify
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
@ -161,6 +162,11 @@ async def _store_document_with_chunks(
if client._config.storage.auto_vacuum:
client._schedule_vacuum()
event = IngestEvent(
documents=[stored_doc],
operation="create" if existing is None else "update",
)
await notify(client._hooks, "after_ingest", client, event)
return stored_doc
@ -212,6 +218,8 @@ async def _update_document_with_chunks(
if client._config.storage.auto_vacuum:
client._schedule_vacuum()
event = IngestEvent(documents=[updated_doc], operation="update")
await notify(client._hooks, "after_ingest", client, event)
return updated_doc
@ -314,6 +322,8 @@ async def _store_documents_with_chunks(
if client._config.storage.auto_vacuum:
client._schedule_vacuum()
event = IngestEvent(documents=created, operation="create")
await notify(client._hooks, "after_ingest", client, event)
return created

View file

@ -35,12 +35,30 @@ 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:
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
# A hook may have cleared it on the way through. after_search reads the
# same request, so it has to end up carrying what retrieval ran with.
request.search_type = request.search_type or "hybrid"
search_type = request.search_type
reranker = client.reranker
@ -58,6 +76,8 @@ async def search(
await _attach_picture_data(client, chunks)
chunk_results = await reranker.rerank(query, chunks, top_n=limit)
else:
# Image queries are vector-only whatever the caller asked for.
request.search_type = "vector"
embedder = client.embedder
if not embedder.supports_images:
raise ValueError(
@ -79,6 +99,9 @@ async def search(
if include_images:
await _populate_image_data(client, results)
for hook in client._hooks:
results = await hook.after_search(client, request, results)
return results

View file

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

View file

@ -394,6 +394,7 @@ def _build_result(
surviving_refs = set(refs)
merged_image_data: dict[str, str] = {}
merged_captions: dict[str, str] = {}
merged_annotations: dict[str, None] = {}
for r in original_results:
if r.doc_item_refs and not surviving_refs.intersection(r.doc_item_refs):
continue
@ -401,6 +402,8 @@ def _build_result(
merged_image_data.update(r.image_data)
if r.picture_captions:
merged_captions.update(r.picture_captions)
if r.annotations:
merged_annotations.update(dict.fromkeys(r.annotations))
return SearchResult(
content=expanded_content,
@ -418,6 +421,7 @@ def _build_result(
labels=sorted(labels) or first.labels,
image_data=merged_image_data or None,
picture_captions=merged_captions,
annotations=list(merged_annotations) or None,
)

View file

@ -0,0 +1,234 @@
import logging
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
from dataclasses import dataclass
from importlib.metadata import entry_points
from types import TracebackType
from typing import TYPE_CHECKING, Literal, Protocol, overload, 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
ENTRY_POINT_GROUP = "haiku.rag.hooks"
logger = logging.getLogger(__name__)
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.
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).
"""
@asynccontextmanager
async def lifespan(self, client: "HaikuRAG") -> AsyncIterator[None]:
"""Hold whatever resources the hook needs for as long as the client is
open. Factories stay resource-free; acquire here instead.
Entered once the store is open, in configured order, and exited in
reverse order while the store, embedder and reranker are still usable.
Raising on entry fails client entry and unwinds the hooks already
started. Raising on exit is logged and swallowed, and a hook can never
suppress an exception raised by the client's caller.
A hook running background work owns stopping it. Cancel or signal the
tasks before whatever awaits them: a clean ``asyncio.TaskGroup`` exit
waits for its children without cancelling them, so an endless task
parked in one hangs shutdown instead of ending it.
"""
yield
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.
Best-effort observer: the operation has already committed, so
exceptions are logged and never raised, and subsequent hooks still
run. Correctness-critical derived state needs its own
reconciliation. Post-commit hooks are not a supported transformation
point: mutating event models does not alter the committed record,
and explicit client writes are separate operations, not atomic with
the original write."""
async def after_delete(self, client: "HaikuRAG", event: DeleteEvent) -> None:
"""``event.documents`` were deleted; cascades arrive as one event.
Best-effort observer with the same contract as ``after_ingest``."""
async def before_search(
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",
request: SearchRequest,
results: list[SearchResult],
) -> list[SearchResult]:
"""Transform or annotate search results before they are returned.
``request`` reflects any ``before_search`` transformations, and its
``search_type`` is the one retrieval ran with: ``"hybrid"`` where a
text search was left unset, ``"vector"`` for an image query."""
return results
def _lifespan_exit(hook: Hook, lifespan: AbstractAsyncContextManager[None]):
"""Wrap a started lifespan's exit so a teardown failure is logged rather
than raised, and so the hook cannot suppress the caller's exception."""
async def _exit(
exc_type: type[BaseException] | None,
exc: BaseException | None,
tb: TracebackType | None,
) -> bool:
try:
await lifespan.__aexit__(exc_type, exc, tb)
except Exception:
cls = type(hook)
logger.exception(
"%s.%s lifespan teardown failed", cls.__module__, cls.__qualname__
)
return False
return _exit
async def enter_lifespans(
stack: AsyncExitStack, hooks: Sequence[Hook], client: "HaikuRAG"
) -> None:
"""Start each hook's lifespan on ``stack`` in order. A hook that fails to
start propagates, leaving the earlier hooks registered on the stack for the
caller to unwind."""
for hook in hooks:
lifespan = hook.lifespan(client)
await lifespan.__aenter__()
stack.push_async_exit(_lifespan_exit(hook, lifespan))
@overload
async def notify(
hooks: Sequence[Hook],
method: Literal["after_ingest"],
client: "HaikuRAG",
event: IngestEvent,
) -> None: ...
@overload
async def notify(
hooks: Sequence[Hook],
method: Literal["after_delete"],
client: "HaikuRAG",
event: DeleteEvent,
) -> None: ...
async def notify(
hooks: Sequence[Hook],
method: Literal["after_ingest", "after_delete"],
client: "HaikuRAG",
event: IngestEvent | DeleteEvent,
) -> None:
"""Fire post-commit observer hooks best-effort: a hook failure is logged
and never raised (the operation already committed), and subsequent hooks
still run."""
for hook in hooks:
try:
await getattr(hook, method)(client, event)
except Exception:
cls = type(hook)
logger.exception(
"%s.%s.%s failed for documents %s",
cls.__module__,
cls.__qualname__,
method,
[d.id for d in event.documents],
)
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

View file

@ -136,6 +136,10 @@ class SearchResult(BaseModel):
``chunk_meta`` is the anchor chunk's unparsed ``Chunk.metadata`` and does not
include the metadata of any other chunks merged with it. Never part of
``format_for_agent`` output.
``annotations`` carries free-text notes attached by ``after_search``
hooks (e.g. definitions of terms appearing in the content). They
survive context expansion and render as notes in ``format_for_agent``.
"""
content: str
@ -154,6 +158,7 @@ class SearchResult(BaseModel):
labels: list[str] = []
image_data: dict[str, str] | None = None
picture_captions: dict[str, str] = {}
annotations: list[str] | None = None
@classmethod
def from_chunk(
@ -225,6 +230,10 @@ class SearchResult(BaseModel):
if caption:
parts.append(f"Figure caption ({self_ref}): {caption}")
if self.annotations:
for note in self.annotations:
parts.append(f"Note: {note}")
# The actual content
parts.append(f"Content:\n{self.content}")

File diff suppressed because one or more lines are too long

View file

@ -1,4 +1,6 @@
import logging
from datetime import UTC, datetime
from pathlib import Path
from unittest.mock import AsyncMock
import httpx
@ -13,10 +15,12 @@ from obstore.exceptions import (
)
from haiku.rag.client import HaikuRAG
from haiku.rag.hooks import Hook
from haiku.rag.ingester.exceptions import PermanentError, TransientError
from haiku.rag.ingester.queue.models import Job, JobOp, JobStatus
from haiku.rag.ingester.workers.pipeline import run_job
from haiku.rag.sources.base import FetchResult, FileTooLargeError, Source
from haiku.rag.sources.fs import FSSource
from haiku.rag.store.models.document import Document
@ -580,3 +584,37 @@ async def test_other_obstore_errors_classified_transient(exc_class):
client.create_document_from_source.side_effect = exc_class("upstream hiccup")
with pytest.raises(TransientError):
await run_job(client, _job())
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent / "cassettes" / "test_pipeline")
class _ThrowingIngestHook(Hook):
async def after_ingest(self, client, event):
raise RuntimeError("ingest hook boom")
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_after_ingest_hook_failure_does_not_fail_job(
temp_db_path, tmp_path, caplog
):
"""A throwing post-commit hook must not turn a committed ingest into a
failed job: run_job returns a successful JobResult instead of raising
into _classify."""
file_path = tmp_path / "doc.md"
file_path.write_text("hello")
fs = FSSource(root=tmp_path, source_id="src")
async with HaikuRAG(temp_db_path, create=True) as client:
client._hooks = [_ThrowingIngestHook()]
with caplog.at_level(logging.ERROR, logger="haiku.rag.hooks"):
result = await run_job(client, _job(uri=str(file_path)), sources=[fs])
assert result.document_id is not None
assert result.deleted is False
assert any("after_ingest" in r.message for r in caplog.records)

717
tests/test_hooks.py Normal file
View file

@ -0,0 +1,717 @@
import logging
from contextlib import asynccontextmanager
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_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, event):
self.events.append(
("ingest", event.operation, tuple((d.id, d.uri) for d in event.documents))
)
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, request):
request.query = f"{request.query} {self.token}"
return request
class FilterHook(Hook):
async def before_search(self, client, request):
request.filter = "uri = 'mem://hooked'"
return request
class ClearSearchTypeHook(Hook):
async def before_search(self, client, request):
request.search_type = None
return request
class SpyAfterSearchHook(Hook):
def __init__(self):
self.search_types: list[str | None] = []
async def after_search(self, client, request, results):
self.search_types.append(request.search_type)
return results
class ReverseResultsHook(Hook):
async def after_search(self, client, request, results):
self.seen_query = request.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 = get_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 = get_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
captured["search_type"] = search_type
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:
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", get_config().search.limit)]
class SpyBeforeSearchHook(Hook):
def __init__(self):
self.requests: list[tuple] = []
async def before_search(self, client, request):
self.requests.append((request.query, request.search_type, request.limit))
return request
@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] * get_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.requests == []
@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 = get_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", "create", ((doc.id, "mem://a"),))]
# A batch import arrives as one event carrying all documents.
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",
"create",
((batch[0].id, "mem://b"), (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", "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
async def test_metadata_only_update_does_not_fire_after_ingest(temp_db_path):
spy = RecordingHook()
dim = get_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 = get_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
# 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, request, results):
for result in results:
result.annotations = ["XMT: transmit"]
return results
@pytest.mark.asyncio
async def test_after_search_hook_annotations_render_for_agent(temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as client:
client._hooks = [AnnotateHook()]
async def fake_search(query, limit, search_type=None, filter=None, **kwargs):
return [(Chunk(id="c1", content="XMT lamp check", order=0), 0.9)]
client.chunk_repository.search = fake_search
results = await client.search("lamp", include_images=False)
assert results[0].annotations == ["XMT: transmit"]
assert "Note: XMT: transmit" in results[0].format_for_agent()
def test_format_for_agent_without_annotations_has_no_notes():
from haiku.rag.store.models.chunk import SearchResult
result = SearchResult(content="plain", score=0.5)
assert "Note:" not in result.format_for_agent()
@pytest.mark.asyncio
async def test_annotations_survive_context_expansion(temp_db_path):
from haiku.rag.context import expand_with_items, window_for
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
async with HaikuRAG(temp_db_path, create=True) as client:
items = [
DocumentItem(
document_id="doc-1",
position=i,
self_ref=f"#/texts/{i}",
label="text",
text=f"Paragraph {i}. " * 10,
)
for i in range(5)
]
await client.document_item_repository.create_items("doc-1", items)
r1 = SearchResult(
content="Paragraph 1.",
score=0.9,
chunk_id="c1",
document_id="doc-1",
doc_item_refs=["#/texts/1"],
annotations=["XMT: transmit", "shared note"],
)
r2 = SearchResult(
content="Paragraph 3.",
score=0.85,
chunk_id="c2",
document_id="doc-1",
doc_item_refs=["#/texts/3"],
annotations=["RCV: receive", "shared note"],
)
repo = client.document_item_repository
positions = (
await repo.resolve_refs_grouped(
{"doc-1": [ref for r in (r1, r2) for ref in r.doc_item_refs]}
)
)["doc-1"]
window_items = (
await repo.get_items_in_ranges({"doc-1": window_for(positions)})
)["doc-1"]
expanded = expand_with_items([r1, r2], 5000, positions, window_items)
assert len(expanded) == 1
assert expanded[0].annotations == [
"XMT: transmit",
"shared note",
"RCV: receive",
]
class ThrowingHook(Hook):
async def after_ingest(self, client, event):
raise RuntimeError("ingest hook boom")
async def after_delete(self, client, event):
raise RuntimeError("delete hook boom")
async def before_search(self, client, request):
raise RuntimeError("search hook boom")
@pytest.mark.asyncio
async def test_after_ingest_hook_failure_is_logged_not_raised(temp_db_path, caplog):
spy = RecordingHook()
dim = get_config().embeddings.model.vector_dim
async with HaikuRAG(temp_db_path, create=True) as client:
client._hooks = [ThrowingHook(), spy]
with caplog.at_level(logging.ERROR, logger="haiku.rag.hooks"):
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
stored = await client.get_document_by_id(doc.id)
assert stored is not None
# Subsequent hooks still run after a failing one.
assert spy.events == [("ingest", "create", ((doc.id, "mem://a"),))]
record = next(r for r in caplog.records if "after_ingest" in r.message)
assert "tests.test_hooks.ThrowingHook" in record.message
assert str(doc.id) in record.message
@pytest.mark.asyncio
async def test_after_delete_hook_failure_is_logged_not_raised(temp_db_path, caplog):
spy = RecordingHook()
dim = get_config().embeddings.model.vector_dim
async with HaikuRAG(temp_db_path, create=True) as client:
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
client._hooks = [ThrowingHook(), spy]
with caplog.at_level(logging.ERROR, logger="haiku.rag.hooks"):
assert await client.delete_document(doc.id) is True
assert await client.get_document_by_id(doc.id) is None
assert spy.events == [("delete", ((doc.id, "mem://a"),))]
record = next(r for r in caplog.records if "after_delete" in r.message)
assert "tests.test_hooks.ThrowingHook" in record.message
assert str(doc.id) in record.message
@pytest.mark.asyncio
async def test_before_search_hook_failure_propagates(temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as client:
client._hooks = [ThrowingHook()]
with pytest.raises(RuntimeError, match="search hook boom"):
await client.search("alpha")
@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 == []
class LifespanHook(Hook):
"""Records lifespan transitions into a shared log, so ordering across
several hooks is observable."""
def __init__(self, name: str, log: list[str]):
self.name = name
self.log = log
@asynccontextmanager
async def lifespan(self, client):
self.log.append(f"enter {self.name}")
try:
yield
finally:
self.log.append(f"exit {self.name}")
class ExceptionRecordingHook(Hook):
"""Records whatever exception its lifespan exit was told about, then lets
it continue on its way."""
def __init__(self, seen: list[str]):
self.seen = seen
@asynccontextmanager
async def lifespan(self, client):
try:
yield
except Exception as exc:
self.seen.append(str(exc))
raise
@pytest.mark.asyncio
async def test_lifespans_enter_in_order_and_exit_in_reverse(temp_db_path):
log: list[str] = []
client = HaikuRAG(temp_db_path, create=True)
client._hooks = [LifespanHook("a", log), LifespanHook("b", log)]
async with client:
assert log == ["enter a", "enter b"]
assert log == ["enter a", "enter b", "exit b", "exit a"]
@pytest.mark.asyncio
async def test_lifespan_can_use_the_store_on_entry_and_exit(temp_db_path):
counts: list[int] = []
class _StoreUsingHook(Hook):
@asynccontextmanager
async def lifespan(self, client):
counts.append(len(await client.list_documents()))
try:
yield
finally:
counts.append(len(await client.list_documents()))
client = HaikuRAG(temp_db_path, create=True)
client._hooks = [_StoreUsingHook()]
dim = get_config().embeddings.model.vector_dim
async with client:
await client.import_document(
_docling_doc("a", "Alpha body"),
[Chunk(content="Alpha body", embedding=[0.1] * dim, order=0)],
uri="mem://lifespan",
title="Alpha",
)
assert counts == [0, 1]
@pytest.mark.asyncio
async def test_lifespan_startup_failure_unwinds_started_hooks(temp_db_path):
log: list[str] = []
class _FailingStartHook(Hook):
@asynccontextmanager
async def lifespan(self, client):
raise RuntimeError("cannot start")
yield # unreachable; asynccontextmanager needs a generator
client = HaikuRAG(temp_db_path, create=True)
client._hooks = [LifespanHook("a", log), _FailingStartHook()]
with pytest.raises(RuntimeError, match="cannot start"):
async with client:
pass
assert log == ["enter a", "exit a"]
assert not client.store.db.is_open()
@pytest.mark.asyncio
async def test_lifespan_teardown_failure_is_logged_and_suppressed(temp_db_path, caplog):
log: list[str] = []
class _FailingExitHook(Hook):
@asynccontextmanager
async def lifespan(self, client):
yield
raise RuntimeError("cannot stop")
client = HaikuRAG(temp_db_path, create=True)
client._hooks = [LifespanHook("a", log), _FailingExitHook()]
with caplog.at_level(logging.ERROR, logger="haiku.rag.hooks"):
async with client:
pass
# The surviving hook still exits, and teardown does not raise.
assert log == ["enter a", "exit a"]
assert any("lifespan" in r.message for r in caplog.records)
@pytest.mark.asyncio
async def test_lifespans_see_the_exception_being_unwound(temp_db_path):
seen: list[str] = []
client = HaikuRAG(temp_db_path, create=True)
client._hooks = [ExceptionRecordingHook(seen)]
with pytest.raises(ValueError, match="from the body"):
async with client:
raise ValueError("from the body")
assert seen == ["from the body"]
@pytest.mark.asyncio
async def test_a_swallowing_lifespan_hides_nothing_from_anyone(temp_db_path):
"""A hook that eats the exception in its own teardown must neither
suppress it for the caller nor make the hooks unwound after it believe
the shutdown was clean."""
seen: list[str] = []
class _SwallowingHook(Hook):
@asynccontextmanager
async def lifespan(self, client):
try:
yield
except Exception:
pass
client = HaikuRAG(temp_db_path, create=True)
# The swallowing hook is entered last, so it unwinds first.
client._hooks = [ExceptionRecordingHook(seen), _SwallowingHook()]
with pytest.raises(ValueError, match="from the body"):
async with client:
raise ValueError("from the body")
assert seen == ["from the body"]
@pytest.mark.asyncio
async def test_default_lifespan_is_a_noop(temp_db_path):
spy = RecordingHook()
client = HaikuRAG(temp_db_path, create=True)
client._hooks = [spy]
dim = get_config().embeddings.model.vector_dim
async with client:
doc = await client.import_document(
_docling_doc("a", "Alpha body"),
[Chunk(content="Alpha body", embedding=[0.1] * dim, order=0)],
uri="mem://default-lifespan",
title="Alpha",
)
# A hook that overrides no lifespan still reaches its other hook points.
assert spy.events == [("ingest", "create", ((doc.id, "mem://default-lifespan"),))]
@pytest.mark.asyncio
async def test_startup_failure_is_forwarded_to_started_lifespans(temp_db_path):
"""A hook that fails to start is an unwind like any other: the lifespans
already running are told what went wrong, not handed a clean shutdown."""
seen: list[str] = []
class _FailingStartHook(Hook):
@asynccontextmanager
async def lifespan(self, client):
raise RuntimeError("cannot start")
yield # unreachable; asynccontextmanager needs a generator
client = HaikuRAG(temp_db_path, create=True)
client._hooks = [ExceptionRecordingHook(seen), _FailingStartHook()]
with pytest.raises(RuntimeError, match="cannot start"):
async with client:
pass
assert seen == ["cannot start"]
@pytest.mark.asyncio
async def test_after_search_sees_the_search_type_that_ran(temp_db_path):
"""A before_search hook may leave search_type unset. Retrieval falls back
to hybrid, so the request after_search reads must say hybrid too."""
async with HaikuRAG(temp_db_path, create=True) as client:
spy = SpyAfterSearchHook()
client._hooks = [ClearSearchTypeHook(), spy]
captured = await _capture_repo_search(client)
await client.search("alpha", include_images=False)
assert captured["search_type"] == "hybrid"
assert spy.search_types == ["hybrid"]
@pytest.mark.asyncio
async def test_after_search_reports_vector_for_image_queries(temp_db_path):
"""Image queries run vector-only whatever the caller asked for, so the
request must not still be advertising the caller's choice."""
async with HaikuRAG(temp_db_path, create=True) as client:
spy = SpyAfterSearchHook()
client._hooks = [spy]
async def fake_search(query, limit, search_type=None, filter=None, **kwargs):
return []
client.chunk_repository.search = fake_search
async def fake_embed_image(image):
return [0.1] * get_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", search_type="fts", include_images=False)
assert spy.search_types == ["vector"]

View file

@ -43,6 +43,7 @@ nav = [
{ Develop = [
{ Python = "python.md" },
{ "Custom pipelines" = "custom-pipelines.md" },
{ Hooks = "hooks.md" },
{ Toolsets = "tools.md" },
{ "Web app" = "apps.md" },
] },