diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ab6d844..5c2bbfc6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### 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. +- `SearchResult.annotations` carries notes attached by `after_search` hooks, preserved through context expansion and rendered in agent-facing output. ## [0.78.0] - 2026-08-24 diff --git a/README.md b/README.md index 1cc489b0..934fd452 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/docs/configuration/index.md b/docs/configuration/index.md index d2434af7..e83c795b 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -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 diff --git a/docs/hooks.md b/docs/hooks.md new file mode 100644 index 00000000..8f894983 --- /dev/null +++ b/docs/hooks.md @@ -0,0 +1,67 @@ +# 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, document)` | A document's content was written (create, import, batch import, update) | Deriving state from documents | +| `after_delete(client, document_id)` | A document was deleted, once per document in a cascade | Cleaning up derived state | +| `before_search(client, query, filter)` | Before retrieval, text queries only | Query expansion, filter injection | +| `after_search(client, query, results)` | After retrieval, reranking, and deduplication | Annotating, reordering, or filtering results | + +`before_search` returns the `(query, filter)` pair to search with. The returned query feeds both the vector and the full-text side. `after_search` returns the result list. Hooks run in the order listed in config, each receiving the previous hook's output. + +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, query, filter): + expanded = my_glossary.expand(query) + return expanded, filter + + async def after_search(self, client, query, 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. + +## 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 + +- **Update equals ingest.** `after_ingest` fires for both creation and content updates. Treat it as "replace any state you derived from this document". Metadata-only and title-only updates do not fire. +- **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. diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index d9785701..03f50d98 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -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, ) diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 2a75fd4a..765282ef 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -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}") diff --git a/tests/test_hooks.py b/tests/test_hooks.py index 45696357..51d74062 100644 --- a/tests/test_hooks.py +++ b/tests/test_hooks.py @@ -265,6 +265,84 @@ async def test_after_delete_fires_for_cascade(temp_db_path): assert all(event[0] == "delete" for event in spy.events) +class AnnotateHook(Hook): + async def after_search(self, client, query, 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 + 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"], + ) + + expanded = await expand_with_items( + client.document_item_repository, "doc-1", [r1, r2], 5000 + ) + + assert len(expanded) == 1 + assert expanded[0].annotations == [ + "XMT: transmit", + "shared note", + "RCV: receive", + ] + + @pytest.mark.asyncio async def test_delete_missing_document_fires_nothing(temp_db_path): spy = RecordingHook() diff --git a/zensical.toml b/zensical.toml index 5d202ee2..d525e63c 100644 --- a/zensical.toml +++ b/zensical.toml @@ -43,6 +43,7 @@ nav = [ { Develop = [ { Python = "python.md" }, { "Custom pipelines" = "custom-pipelines.md" }, + { Hooks = "hooks.md" }, { Toolsets = "tools.md" }, { "Web app" = "apps.md" }, ] },