09a7076b7e
Comments and docstrings across the branch narrated rejected alternatives, consequences and history; each now states the current contract. Renames test_a_legacy_uri_client_keeps_its_error to test_an_unnamed_database_keeps_its_error. Documents the Sandbox connection paths, the citation header's database segment, both AmbiguousDatabaseError conditions on create_app, and run_inspector's scope parameter. Doc paragraphs added by the branch in python.md, storage.md and cli.md are one physical line each.
95 lines
3.4 KiB
Python
95 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."""
|
|
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]
|