274 lines
8.4 KiB
Python
274 lines
8.4 KiB
Python
import pytest
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import 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, document):
|
|
self.events.append(("ingest", document.id, document.uri))
|
|
|
|
async def after_delete(self, client, document_id):
|
|
self.events.append(("delete", document_id))
|
|
|
|
|
|
class AppendTokenHook(Hook):
|
|
def __init__(self, token: str = "expanded"):
|
|
self.token = token
|
|
|
|
async def before_search(self, client, query, filter):
|
|
return f"{query} {self.token}", filter
|
|
|
|
|
|
class FilterHook(Hook):
|
|
async def before_search(self, client, query, filter):
|
|
return query, "uri = 'mem://hooked'"
|
|
|
|
|
|
class ReverseResultsHook(Hook):
|
|
async def after_search(self, client, query, results):
|
|
self.seen_query = 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 = 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 = 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
|
|
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:
|
|
client._hooks = [AppendTokenHook("one"), AppendTokenHook("two"), FilterHook()]
|
|
captured = await _capture_repo_search(client)
|
|
|
|
await client.search("alpha")
|
|
|
|
assert captured["query"] == "alpha one two"
|
|
assert captured["filter"] == "uri = 'mem://hooked'"
|
|
|
|
|
|
class SpyBeforeSearchHook(Hook):
|
|
def __init__(self):
|
|
self.called: list[str] = []
|
|
|
|
async def before_search(self, client, query, filter):
|
|
self.called.append(query)
|
|
return query, filter
|
|
|
|
|
|
@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] * 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.called == []
|
|
|
|
|
|
@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 = 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", doc.id, "mem://a")]
|
|
|
|
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", batch[0].id, "mem://b"),
|
|
("ingest", 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", doc.id, "mem://a")]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_metadata_only_update_does_not_fire_after_ingest(temp_db_path):
|
|
spy = RecordingHook()
|
|
dim = 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 = 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
|
|
|
|
deleted = {event[1] for event in spy.events}
|
|
assert deleted == {parent.id, child.id}
|
|
assert all(event[0] == "delete" for event in spy.events)
|
|
|
|
|
|
@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 == []
|