ca2e28559e
Two files of 1,405 and 814 lines become seven: scope resolution, lifecycle, search, documents, expansion, citations and capabilities. `_config`, `_seed` and the rest move to `helpers.py`, importable by the sandbox tests that share them, and the package points VCR back at `tests/cassettes/multi_db/`.
96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
"""Databases to run the multi-database tests against."""
|
|
|
|
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.config import get_config
|
|
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
|
from haiku.rag.store.models import Chunk
|
|
from haiku.rag.utils import locate_database
|
|
|
|
|
|
def _config(tmp_path, names) -> AppConfig:
|
|
return AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
databases={n: str(tmp_path / f"{n}.lancedb") for n in names}
|
|
)
|
|
)
|
|
|
|
|
|
async def _seed(config, name, contents):
|
|
"""Precomputed embeddings and FTS queries keep the embedder out of the way:
|
|
these tests are about fusion, not retrieval quality."""
|
|
dim = get_config().embeddings.model.vector_dim
|
|
async with HaikuRAG(config=config, create=True, sources=[name]) as rag:
|
|
for content in contents:
|
|
doc = DoclingDocument(name=content)
|
|
doc.add_text(label=DocItemLabel.TEXT, text=content)
|
|
await rag.import_document(
|
|
doc,
|
|
[Chunk(content=content, embedding=[0.1] * dim, order=0)],
|
|
uri=f"test://{name}/{content}",
|
|
)
|
|
|
|
|
|
async def _restore_embedder(config, name, *, provider=None, model_name=None):
|
|
"""Rewrite what one database records about the embedder that wrote it,
|
|
standing in for a database built elsewhere with another model."""
|
|
import json
|
|
|
|
import lancedb
|
|
|
|
_, db_path = locate_database(config.lancedb.databases[name])
|
|
assert db_path is not None
|
|
db = await lancedb.connect_async(str(db_path.resolve()))
|
|
table = await db.open_table("settings")
|
|
rows = (
|
|
await table.query().where("id = 'settings'").limit(1).to_arrow()
|
|
).to_pylist()
|
|
stored = json.loads(rows[0]["settings"])
|
|
model = stored["embeddings"]["model"]
|
|
if provider is not None:
|
|
model["provider"] = provider
|
|
if model_name is not None:
|
|
model["name"] = model_name
|
|
await table.update({"settings": json.dumps(stored)}, where="id = 'settings'")
|
|
|
|
|
|
async def _seed_expandable(config, name, sentences):
|
|
"""One document whose chunk covers a single item, so expansion has
|
|
neighbours to pull in and rebuilds the result rather than passing it
|
|
through."""
|
|
dim = get_config().embeddings.model.vector_dim
|
|
doc = DoclingDocument(name=name)
|
|
for sentence in sentences:
|
|
doc.add_text(label=DocItemLabel.TEXT, text=sentence)
|
|
async with HaikuRAG(config=config, create=True, sources=[name]) as rag:
|
|
await rag.import_document(
|
|
doc,
|
|
[
|
|
Chunk(
|
|
content=sentences[0],
|
|
embedding=[0.1] * dim,
|
|
order=0,
|
|
metadata={"doc_item_refs": ["#/texts/0"]},
|
|
)
|
|
],
|
|
uri=f"test://{name}/expandable",
|
|
)
|
|
|
|
|
|
class StubReranker:
|
|
"""Scores the union, reversing it so the ordering is unmistakably its own."""
|
|
|
|
def __init__(self):
|
|
self.seen: list[str] = []
|
|
|
|
async def rerank(self, query, chunks, top_n):
|
|
self.seen = [c.content for c in chunks]
|
|
# Whatever the caller attached before handing them over.
|
|
self.attached = {
|
|
c.content.split()[0]: c._picture_data
|
|
for c in chunks
|
|
if getattr(c, "_picture_data", None)
|
|
}
|
|
return [(c, 1.0 - i) for i, c in enumerate(reversed(chunks))][:top_n]
|