Hooks own resources through a lifespan
This commit is contained in:
Yiorgis Gozadinos
parent
1cb5e296b0
commit
1ecf77cad5
6 changed files with 410 additions and 29 deletions
|
|
@ -5,6 +5,7 @@
|
|||
|
||||
- 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
|
||||
|
||||
|
|
|
|||
|
|
@ -14,8 +14,9 @@ Subclass `haiku.rag.hooks.Hook` and override any subset:
|
|||
| `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. Hooks run in the order listed in config, each receiving the previous hook's output.
|
||||
Events are batch shaped. `IngestEvent` carries `documents` (a batch import arrives as one event with the whole batch) and `operation` (`"create"` or `"update"`). `DeleteEvent` carries the deleted `documents` in their last-known state, since the rows are already gone; a cascade arrives as one event. `SearchRequest` carries `query`, `filter`, `search_type`, and `limit`. `before_search` returns the request to search with, and may modify any of its fields. The query feeds both the vector and the full-text side. `after_search` returns the result list, 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.
|
||||
|
||||
|
|
@ -54,6 +55,46 @@ hooks:
|
|||
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -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,7 +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, load_hooks, notify
|
||||
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
|
||||
|
|
@ -118,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")
|
||||
|
|
@ -152,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.
|
||||
|
|
|
|||
|
|
@ -55,7 +55,10 @@ async def search(
|
|||
assert isinstance(query, str), "before_search must keep text queries text"
|
||||
filter = request.filter
|
||||
limit = request.limit
|
||||
search_type = request.search_type or "hybrid"
|
||||
# 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
|
||||
|
||||
|
|
@ -73,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(
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
import logging
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
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
|
||||
|
|
@ -59,6 +61,24 @@ class 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
|
||||
|
|
@ -94,10 +114,45 @@ class Hook:
|
|||
results: list[SearchResult],
|
||||
) -> list[SearchResult]:
|
||||
"""Transform or annotate search results before they are returned.
|
||||
``request`` reflects any ``before_search`` transformations."""
|
||||
``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],
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -37,6 +38,21 @@ class FilterHook(Hook):
|
|||
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
|
||||
|
|
@ -108,6 +124,7 @@ async def _capture_repo_search(client):
|
|||
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
|
||||
|
|
@ -462,3 +479,239 @@ async def test_delete_missing_document_fires_nothing(temp_db_path):
|
|||
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"]
|
||||
|
|
|
|||
Loading…
Reference in a new issue