Compare commits
6 commits
main
...
feat/hooks
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1ecf77cad5 | ||
|
|
1cb5e296b0 | ||
|
|
e10762854c | ||
|
|
8083c57246 | ||
|
|
b4a3c678a0 | ||
|
|
330468f2ca |
15 changed files with 1263 additions and 31 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
111
docs/hooks.md
Normal 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.
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
234
haiku_rag_slim/haiku/rag/hooks.py
Normal file
234
haiku_rag_slim/haiku/rag/hooks.py
Normal 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
|
||||
|
|
@ -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
|
|
@ -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
717
tests/test_hooks.py
Normal 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"]
|
||||
|
|
@ -43,6 +43,7 @@ nav = [
|
|||
{ Develop = [
|
||||
{ Python = "python.md" },
|
||||
{ "Custom pipelines" = "custom-pipelines.md" },
|
||||
{ Hooks = "hooks.md" },
|
||||
{ Toolsets = "tools.md" },
|
||||
{ "Web app" = "apps.md" },
|
||||
] },
|
||||
|
|
|
|||
Loading…
Reference in a new issue