russell@unturf.com
a4e1dc9a10
A stable façade so another Python app can use arborist as a content-addressed / Merkle / audit-chained store without the CLI or a wire protocol. Import from arborist.embed, not internal modules, so refactors don't break embedders. Surface: open_store(path), ingest_documents(conn, docs), search(conn, q), plus re-exported Document/Edge/Source/Hit/IngestStats. Core only (python+sqlite3) — no extras. _IterableSource adapts a plain doc iterable into the Source contract. This is the seam for using arborist as neopig's optional provenance backend: neopig produces Documents from crawled pages, arborist gives content-dedup (document_root) + FTS5 + an append-only audit chain alongside neopig's existing md5/FileVault storage. Docs in docs/embedding.md. 6 tests pin open/ingest/dedup/idempotence/edges/search.
111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
"""Tests for arborist.embed — the supported embedding surface.
|
|
|
|
These pin the contract a downstream Python app (e.g. neopig) depends on:
|
|
open a store, ingest Documents, get content-level dedup, search them.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from arborist.embed import (
|
|
Document,
|
|
Edge,
|
|
Source,
|
|
open_store,
|
|
ingest_documents,
|
|
search,
|
|
)
|
|
|
|
|
|
def _doc(uri, content, title=None, source_type="embedded"):
|
|
return Document(uri=uri, content=content, source_type=source_type, title=title)
|
|
|
|
|
|
def test_open_store_creates_schema(tmp_path):
|
|
conn = open_store(tmp_path / "embed.db")
|
|
try:
|
|
# A fresh store has the documents table (schema migrated on open).
|
|
names = {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")}
|
|
assert "documents" in names
|
|
assert "audit_events" in names
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def test_ingest_documents_lands_and_is_searchable(tmp_path):
|
|
conn = open_store(tmp_path / "embed.db")
|
|
try:
|
|
stats = ingest_documents(
|
|
conn,
|
|
[
|
|
_doc("https://x.com/a", "the quick brown fox jumps", title="A"),
|
|
_doc("https://x.com/b", "lorem ipsum dolor sit amet", title="B"),
|
|
],
|
|
)
|
|
assert stats.inserted == 2
|
|
n = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
|
|
assert n == 2
|
|
|
|
hits = search(conn, "brown fox", limit=10)
|
|
assert any("x.com/a" in h.document_uri for h in hits), hits
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def test_ingest_is_idempotent_on_content(tmp_path):
|
|
"""Same content → same document_root → no-op re-insert (dedup)."""
|
|
conn = open_store(tmp_path / "embed.db")
|
|
try:
|
|
ingest_documents(conn, [_doc("https://x.com/a", "identical body text here")])
|
|
first = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
|
|
# Re-ingest the same content (same uri) — must not duplicate.
|
|
stats = ingest_documents(conn, [_doc("https://x.com/a", "identical body text here")])
|
|
second = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
|
|
assert first == second == 1
|
|
assert stats.inserted == 0
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def test_dedup_is_content_not_uri(tmp_path):
|
|
"""Two different URIs with identical content share one document_root —
|
|
the property an embedder leans on for cross-source dedup."""
|
|
conn = open_store(tmp_path / "embed.db")
|
|
try:
|
|
ingest_documents(
|
|
conn,
|
|
[
|
|
_doc("https://x.com/a", "byte identical content"),
|
|
_doc("https://y.com/mirror", "byte identical content"),
|
|
],
|
|
)
|
|
roots = [r[0] for r in conn.execute("SELECT DISTINCT document_root FROM documents")]
|
|
assert len(roots) == 1, "identical content must collapse to one root"
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def test_edges_survive_ingest(tmp_path):
|
|
conn = open_store(tmp_path / "embed.db")
|
|
try:
|
|
page = Document(
|
|
uri="https://x.com/page",
|
|
content="a page that links to an image",
|
|
source_type="neopig_html",
|
|
title="Page",
|
|
edges=[Edge(edge_type="embeds_media", dst_uri="https://x.com/img.jpg")],
|
|
)
|
|
ingest_documents(conn, [page])
|
|
n_edges = conn.execute("SELECT COUNT(*) FROM edges").fetchone()[0]
|
|
assert n_edges >= 1
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def test_iterable_source_is_a_source_subclass():
|
|
"""The embed wrapper must satisfy the Source contract so it flows
|
|
through the same ingest pipeline as any first-class corpus."""
|
|
from arborist.embed import _IterableSource
|
|
|
|
src = _IterableSource([], source_type="embedded")
|
|
assert isinstance(src, Source)
|
|
assert src.source_type == "embedded"
|