202cbd2d3f
Reads tagged the document they returned and writes did not, so creating in `alpha` came back with `source=None` while reading the same document came back with `"alpha"`. Every outward write result goes through the owning session.
334 lines
13 KiB
Python
334 lines
13 KiB
Python
"""Listing and looking up documents across databases."""
|
|
|
|
import pytest
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.client.session import SingleDatabaseSession
|
|
from haiku.rag.config import get_config
|
|
from haiku.rag.store.models import Chunk
|
|
from tests.multi_db.helpers import (
|
|
_config,
|
|
_seed,
|
|
)
|
|
|
|
|
|
class TestListingAcrossDatabases:
|
|
"""The chat TUI's document filter lists documents through the client, and a
|
|
client covering a set has no repositories of its own."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_listing_covers_every_database(self, tmp_path):
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one", "alpha two"])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
docs = await rag.list_documents()
|
|
|
|
assert {d.uri for d in docs} == {
|
|
"test://alpha/alpha one",
|
|
"test://alpha/alpha two",
|
|
"test://beta/beta one",
|
|
}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_counting_covers_every_database(self, tmp_path):
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one", "alpha two"])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
assert await rag.count_documents() == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_limit_bounds_the_merged_listing(self, tmp_path):
|
|
"""A limit is that many documents in total, not that many per database."""
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one", "alpha two"])
|
|
await _seed(config, "beta", ["beta one", "beta two"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
assert len(await rag.list_documents(limit=3)) == 3
|
|
assert len(await rag.list_documents(limit=2, offset=2)) == 2
|
|
assert len(await rag.list_documents(offset=3)) == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_page_shows_every_database(self, tmp_path):
|
|
"""A window is taken across the databases, not filled from the first one:
|
|
concatenating hides every database after whichever was listed first."""
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", [f"alpha {i}" for i in range(5)])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
page = await rag.list_documents(limit=3)
|
|
|
|
assert len(page) == 3
|
|
assert {(d.uri or "").split("/")[2] for d in page} == {"alpha", "beta"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_filter_reaches_every_database(self, tmp_path):
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one"])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
docs = await rag.list_documents(filter="uri LIKE 'test://beta/%'")
|
|
|
|
assert [d.uri for d in docs] == ["test://beta/beta one"]
|
|
|
|
|
|
class TestLookupByIdentifier:
|
|
"""An id or a URI says nothing about which database holds it, and a client
|
|
covering a set has no repositories of its own."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_document_is_found_in_whichever_database_holds_it(self, tmp_path):
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one"])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
beta = (await rag.clients_for(["beta"]))[0]
|
|
[target] = await beta.document_repository.list_all(limit=1)
|
|
assert target.id is not None
|
|
|
|
found = await rag.get_document_by_id(target.id)
|
|
by_uri = await rag.get_document_by_uri("test://alpha/alpha one")
|
|
resolved = await rag.resolve_document(target.id)
|
|
|
|
assert found is not None and found.uri == "test://beta/beta one"
|
|
assert by_uri is not None and by_uri.uri == "test://alpha/alpha one"
|
|
assert resolved is not None and resolved.uri == "test://beta/beta one"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_chunk_is_found_in_whichever_database_holds_it(self, tmp_path):
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one"])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
beta = (await rag.clients_for(["beta"]))[0]
|
|
[chunk] = await beta.chunk_repository.list_all(limit=1)
|
|
assert chunk.id is not None
|
|
|
|
found = await rag.get_chunk_by_id(chunk.id)
|
|
|
|
assert found is not None and found.content == "beta one"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_document_held_by_two_databases_answers_from_the_first(
|
|
self, tmp_path
|
|
):
|
|
"""A database copied from another holds the same ids. A read has an
|
|
answer wherever it finds one, and which one it is has to be the
|
|
configured order rather than whichever replied first."""
|
|
import shutil
|
|
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one"])
|
|
shutil.copytree(tmp_path / "alpha.lancedb", tmp_path / "beta.lancedb")
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
beta = (await rag.clients_for(["beta"]))[0]
|
|
[target] = await beta.document_repository.list_all(limit=1)
|
|
assert target.id is not None
|
|
|
|
found = await rag.get_document_by_id(target.id)
|
|
|
|
assert found is not None and found.source == "alpha"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_an_unknown_identifier_is_absent_rather_than_an_error(self, tmp_path):
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one"])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
assert (
|
|
await rag.get_document_by_id("00000000-0000-4000-8000-000000000000")
|
|
is None
|
|
)
|
|
assert (
|
|
await rag.get_chunk_by_id("00000000-0000-4000-8000-000000000000")
|
|
is None
|
|
)
|
|
assert await rag.get_document_by_uri("test://nowhere") is None
|
|
|
|
|
|
class TestWritesNameTheirDatabase:
|
|
"""A write returns the document it wrote, and it came from a database. A
|
|
read of the same document names it, so the write has to as well."""
|
|
|
|
@staticmethod
|
|
def _doc(text: str):
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
doc = DoclingDocument(name=text)
|
|
doc.add_text(label=DocItemLabel.TEXT, text=text)
|
|
return doc
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_import_names_the_database(self, tmp_path):
|
|
config = _config(tmp_path, ["alpha"])
|
|
dim = get_config().embeddings.model.vector_dim
|
|
|
|
async with HaikuRAG(config=config, create=True, sources=["alpha"]) as rag:
|
|
written = await rag.import_document(
|
|
self._doc("cats"),
|
|
[Chunk(content="cats", embedding=[0.1] * dim, order=0)],
|
|
uri="test://alpha/cats",
|
|
)
|
|
assert written.id is not None
|
|
read = await rag.get_document_by_id(written.id)
|
|
|
|
assert written.source == "alpha"
|
|
assert read is not None and read.source == written.source
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_batch_import_names_every_document(self, tmp_path):
|
|
from haiku.rag.client.documents import DocumentImport
|
|
|
|
config = _config(tmp_path, ["alpha"])
|
|
dim = get_config().embeddings.model.vector_dim
|
|
|
|
async with HaikuRAG(config=config, create=True, sources=["alpha"]) as rag:
|
|
written = await rag.import_documents(
|
|
[
|
|
DocumentImport(
|
|
docling_document=self._doc(text),
|
|
chunks=[Chunk(content=text, embedding=[0.1] * dim, order=0)],
|
|
uri=f"test://alpha/{text}",
|
|
)
|
|
for text in ("cats", "dogs")
|
|
]
|
|
)
|
|
|
|
assert [d.source for d in written] == ["alpha", "alpha"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_metadata_only_update_names_the_database(self, tmp_path):
|
|
"""Changing only metadata rewrites the row without re-chunking, so it
|
|
never reaches the paths that name a document on the way through."""
|
|
config = _config(tmp_path, ["alpha"])
|
|
dim = get_config().embeddings.model.vector_dim
|
|
|
|
async with HaikuRAG(config=config, create=True, sources=["alpha"]) as rag:
|
|
stored = await rag.import_document(
|
|
self._doc("cats"),
|
|
[Chunk(content="cats", embedding=[0.1] * dim, order=0)],
|
|
uri="test://alpha/cats",
|
|
)
|
|
assert stored.id is not None
|
|
|
|
updated = await rag.update_document(
|
|
stored.id, title="Cats", metadata={"k": "v"}
|
|
)
|
|
|
|
assert updated is not None
|
|
assert updated.title == "Cats"
|
|
assert updated.source == "alpha"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_the_revision_short_circuit_names_the_database(self, tmp_path):
|
|
"""`create_document_from_source` refreshes metadata in place when the
|
|
revision is unchanged, returning the document it rewrote."""
|
|
from haiku.rag.client.documents import _refresh_doc_metadata
|
|
|
|
config = _config(tmp_path, ["alpha"])
|
|
dim = get_config().embeddings.model.vector_dim
|
|
|
|
async with HaikuRAG(config=config, create=True, sources=["alpha"]) as rag:
|
|
stored = await rag.import_document(
|
|
self._doc("cats"),
|
|
[Chunk(content="cats", embedding=[0.1] * dim, order=0)],
|
|
uri="test://alpha/cats",
|
|
)
|
|
assert isinstance(rag._session, SingleDatabaseSession)
|
|
refreshed = await _refresh_doc_metadata(
|
|
rag._session,
|
|
stored,
|
|
title="Cats",
|
|
user_metadata={"k": "v"},
|
|
source_metadata=None,
|
|
)
|
|
|
|
assert refreshed.source == "alpha"
|
|
|
|
|
|
class TestDocumentsNameTheirDatabase:
|
|
"""A listing that spans databases is unreadable when the documents do not
|
|
say which one they came from, the same reason a search result carries one."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_listing_names_each_document_s_database(self, tmp_path):
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one", "alpha two"])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
docs = await rag.list_documents()
|
|
|
|
assert {d.uri: d.source for d in docs} == {
|
|
"test://alpha/alpha one": "alpha",
|
|
"test://alpha/alpha two": "alpha",
|
|
"test://beta/beta one": "beta",
|
|
}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_looked_up_document_names_its_database(self, tmp_path):
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one"])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
beta = (await rag.clients_for(["beta"]))[0]
|
|
[target] = await beta.document_repository.list_all(limit=1)
|
|
assert target.id is not None
|
|
|
|
by_id = await rag.get_document_by_id(target.id)
|
|
by_uri = await rag.get_document_by_uri("test://alpha/alpha one")
|
|
resolved = await rag.resolve_document(target.id)
|
|
|
|
assert by_id is not None and by_id.source == "beta"
|
|
assert by_uri is not None and by_uri.source == "alpha"
|
|
assert resolved is not None and resolved.source == "beta"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_named_database_still_names_itself(self, tmp_path):
|
|
"""`haiku-rag --db-name alpha list` opens one database, and its name is
|
|
the whole reason the option exists."""
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one"])
|
|
await _seed(config, "beta", ["beta one"])
|
|
|
|
async with HaikuRAG(config=config, sources=["alpha"]) as rag:
|
|
[listed] = await rag.list_documents()
|
|
assert listed.id is not None
|
|
by_id = await rag.get_document_by_id(listed.id)
|
|
by_uri = await rag.get_document_by_uri("test://alpha/alpha one")
|
|
|
|
assert listed.source == "alpha"
|
|
assert by_id is not None and by_id.source == "alpha"
|
|
assert by_uri is not None and by_uri.source == "alpha"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_database_leaves_the_source_unset(self, tmp_path, temp_db_path):
|
|
"""Nothing names the database when there is only one to name."""
|
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
|
dim = get_config().embeddings.model.vector_dim
|
|
doc = DoclingDocument(name="solo")
|
|
doc.add_text(label=DocItemLabel.TEXT, text="solo")
|
|
await rag.import_document(
|
|
doc,
|
|
[Chunk(content="solo", embedding=[0.1] * dim, order=0)],
|
|
uri="test://solo",
|
|
)
|
|
|
|
[listed] = await rag.list_documents()
|
|
assert listed.source is None
|
|
assert listed.id is not None
|
|
assert (await rag.get_document_by_id(listed.id)).source is None
|