From 5a731220bf42de808b25e2568e39e6ee84943e45 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Sat, 24 Oct 2026 12:00:28 +0900 Subject: [PATCH] Add example 29: hybrid agent-memory retrieval in one SQL statement Tenant pre-filtered vector.neighbors fused with a full-text match by vector.fuse (RRF), then graph hops kept inside the tenant. The where form needs 26.11.1; the filter form runs on 26.10.2. Test runs both forms. Co-Authored-By: Claude Opus 5.5 --- examples/29_hybrid_memory_retrieval.py | 365 +++++++++++++++++++++++++ examples/README.md | 20 ++ tests/test_example29_hybrid_memory.py | 74 +++++ 3 files changed, 459 insertions(+) create mode 100644 examples/29_hybrid_memory_retrieval.py create mode 100644 tests/test_example29_hybrid_memory.py diff --git a/examples/29_hybrid_memory_retrieval.py b/examples/29_hybrid_memory_retrieval.py new file mode 100644 index 0000000..0c3fac4 --- /dev/null +++ b/examples/29_hybrid_memory_retrieval.py @@ -0,0 +1,365 @@ +#!/usr/bin/env python3 +"""Example 29: Hybrid Agent-Memory Retrieval in One SQL Statement. + +For an agent-memory or GraphRAG builder who today glues a vector store, a +keyword index and a graph store together. This example does the retrieval in +ONE ArcadeDB statement: the nearest facts by embedding, restricted to one +tenant BEFORE the top-k is taken, fused with the full-text matches of the same +tenant by reciprocal rank fusion, and then expanded by graph hops to the facts +that share a related entity. + +Everything goes through db.query / db.command; there is no Python wrapper. +The memory is synthetic (fixed-seed numpy embeddings, templated sentences); +no download. + +The pre-filter is the point. Taking the global top-k first and filtering by +tenant afterwards returns fewer rows than asked, or none, whenever the nearest +neighbours belong to another tenant (this is what Cypher's +``db.index.vector.queryNodes(index, k, vector)`` does: it has no filter, so +``WHERE node.tenant = ...`` runs on its k results; ArcadeData/arcadedb#10685). +``vector.neighbors`` takes the restriction as an option, so the index search +itself only considers the tenant's vectors: + +* ``{ where: { tenant: :tenant } }``: tested on each candidate as the search + meets it. Needs ArcadeDB 26.11.1 (ArcadeData/arcadedb#10169, #10185). This + is the form the statement below uses. +* ``{ filter: (SELECT @rid FROM Fact WHERE tenant = :tenant) }``: a list of + RIDs. Works on 26.10.2 as well. ``--form filter`` runs the example with it, + and the example checks that both forms return the same rows when the engine + has ``where``. + +The index types are created with SQL and are always the same: an ``LSM_VECTOR`` +index on ``Fact.embedding`` and a ``FULL_TEXT`` index on ``Fact.text``. +``vector.fuse`` and the ``SEARCH_INDEX`` source are used as documented in the +ArcadeDB vector-search concept page ("Hybrid Search with vector.fuse"). + +Note on recall: at the index defaults the search beam narrows once an index +holds 10,000 vectors or more (ArcadeData/arcadedb#10704); this memory has +about 900 vectors, far below that, so the defaults are fine here. For a larger +memory pass ``efSearch`` in the options map: ``{ where: ..., efSearch: 100 }``. +""" + +from __future__ import annotations + +import argparse +import re +import shutil +from pathlib import Path + +import numpy as np + +import arcadedb_embedded as arcadedb + +SEED = 20261024 +DIM = 16 +FACTS_PER_TENANT = 150 +TENANTS = ["acme", "globex", "initech", "umbrella", "hooli", "stark"] +REQUESTED_TENANT = "acme" +K = 10 # neighbours the vector side asks for; also the size of the full-text side +TOP = 5 # rows of the fused answer + +# topic -> (entities, sentence templates) +TOPICS = { + "coffee": ( + ["espresso", "grinder", "barista", "beans"], + [ + "drinks {a} every morning and keeps a {b} at home", + "asked the team to order more {a} because the {b} ran out", + ], + ), + "running": ( + ["marathon", "shoes", "pace", "trail"], + [ + "trains for the {a} and tracks the {b} on every run", + "bought new {a} after the {b} race last spring", + ], + ), + "travel": ( + ["lisbon", "flight", "hotel", "passport"], + [ + "plans a trip to {a} and booked the {b} already", + "lost the {a} on the way and had to rebook the {b}", + ], + ), + "music": ( + ["guitar", "piano", "concert", "vinyl"], + [ + "practises the {a} before the {b} on Friday", + "collects {a} records and goes to every {b} in town", + ], + ), +} + +# The one statement. The vector side is pre-filtered by tenant (the `where` option, +# 26.11.1), the full-text side is a subquery on the same tenant, `vector.fuse` merges +# the two ranked lists by reciprocal rank fusion, and the outer SELECT hops through the +# graph from every fused fact: +# related_entities: Fact -MENTIONS-> Entity -RELATES-> Entity (two hops) +# related_fids: ... -MENTIONS<- Fact, without the fact itself (three hops, a +# sample of 3, in no particular order) +# `[tenant = :tenant]` keeps a hop from leaving the tenant: the memory holds one stray +# RELATES edge from an acme entity to a globex entity, and without it the hops of the +# acme facts would return globex facts. +# One statement was possible; nothing is split into several queries. +_HYBRID_TEMPLATE = """ +SELECT record.fid AS fid, + record.text AS fact, + score, + record.out('MENTIONS').out('RELATES')[tenant = :tenant].name + AS related_entities, + difference( + record.out('MENTIONS').out('RELATES').in('MENTIONS')[tenant = :tenant].fid, + [record.fid] + )[0..3] AS related_fids +FROM ( + SELECT expand(`vector.fuse`( + `vector.neighbors`('Fact[embedding]', :qvec, :k, { %(restrict)s }), + (SELECT @rid, $score FROM Fact + WHERE tenant = :tenant AND SEARCH_INDEX('Fact[text]', :keywords) = true + ORDER BY $score DESC LIMIT :k), + { fusion: 'RRF' } + )) +) +ORDER BY score DESC +LIMIT :top +""" +RESTRICT = { + "where": "where: { tenant: :tenant }", # 26.11.1 + "filter": "filter: (SELECT @rid FROM Fact WHERE tenant = :tenant)", # 26.10.2 too +} +HYBRID_SQL = {form: _HYBRID_TEMPLATE % {"restrict": r} for form, r in RESTRICT.items()} + + +def neighbors_sql(restrict: str | None) -> str: + """The vector side alone, with (restrict) or without a restriction.""" + opts = f", {{ {restrict} }}" if restrict else "" + return ( + "SELECT record.fid AS fid, record.tenant AS tenant, distance FROM (" + f"SELECT expand(`vector.neighbors`('Fact[embedding]', :qvec, :k{opts})))" + ) + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--db-path", default="./my_test_databases/hybrid_memory") + p.add_argument( + "--form", + choices=sorted(RESTRICT), + default="where", + help="how the vector side is pre-filtered: where (26.11.1) or filter " + "(a RID subquery, 26.10.2 and later)", + ) + return p.parse_args() + + +def build(db) -> dict: + """Create the schema and a synthetic multi-tenant memory. Returns the rng-made + topic centroids so the caller can write a query vector in the same space.""" + rng = np.random.default_rng(SEED) + centroids = {t: _unit(rng.standard_normal(DIM)) for t in TOPICS} + for ddl in ( + "CREATE VERTEX TYPE Fact", + "CREATE PROPERTY Fact.fid INTEGER", + "CREATE PROPERTY Fact.tenant STRING", + "CREATE PROPERTY Fact.text STRING", + "CREATE PROPERTY Fact.embedding ARRAY_OF_FLOATS", + "CREATE VERTEX TYPE Entity", + "CREATE PROPERTY Entity.eid INTEGER", + "CREATE PROPERTY Entity.tenant STRING", + "CREATE PROPERTY Entity.name STRING", + "CREATE EDGE TYPE MENTIONS", + "CREATE EDGE TYPE RELATES", + "CREATE INDEX ON Fact (fid) UNIQUE_HASH", + "CREATE INDEX ON Entity (eid) UNIQUE_HASH", + "CREATE INDEX ON Fact (text) FULL_TEXT", + 'CREATE INDEX ON Fact (embedding) LSM_VECTOR METADATA {"dimensions": %d, ' + '"similarity": "COSINE"}' % DIM, + ): + db.command("sql", ddl) + + eid = fid = 0 + entity_ids = {} # (tenant, topic, name) -> eid + entity_edges = [] + with db.transaction(): + for tenant in TENANTS: + for topic, (names, _) in TOPICS.items(): + for name in names: + db.command( + "sql", + "CREATE VERTEX Entity SET eid = ?, tenant = ?, name = ?", + eid, + tenant, + name, + ) + entity_ids[(tenant, topic, name)] = eid + eid += 1 + # RELATES: each entity points at the next one of its topic, and the + # last entity of a topic points at the first of the next topic. + topics = list(TOPICS) + for ti, topic in enumerate(topics): + names = TOPICS[topic][0] + for ni, name in enumerate(names): + if ni + 1 < len(names): + nxt = (tenant, topic, names[ni + 1]) + else: + other = topics[(ti + 1) % len(topics)] + nxt = (tenant, other, TOPICS[other][0][0]) + entity_edges.append( + (entity_ids[(tenant, topic, name)], entity_ids[nxt]) + ) + # the stray edge that crosses tenants (see HYBRID_SQL) + entity_edges.append( + ( + entity_ids[("acme", "coffee", "grinder")], + entity_ids[("globex", "music", "guitar")], + ) + ) + for src, dst in entity_edges: + db.command( + "sql", + "CREATE EDGE RELATES FROM (SELECT FROM Entity WHERE eid = ?) " + "TO (SELECT FROM Entity WHERE eid = ?)", + src, + dst, + ) + + mention_edges = [] + with db.transaction(): + for tenant in TENANTS: + # acme's facts are looser around the topic than everyone else's, so + # a query near a topic has other tenants' facts as its nearest + # neighbours: the case where filtering after the top-k loses rows. + noise = 0.45 if tenant == REQUESTED_TENANT else 0.15 + for i in range(FACTS_PER_TENANT): + topic = list(TOPICS)[int(rng.integers(len(TOPICS)))] + names, templates = TOPICS[topic] + a, b = rng.choice(names, size=2, replace=False) + template = templates[int(rng.integers(len(templates)))] + text = f"{tenant} user " + template.format(a=a, b=b) + emb = _unit(centroids[topic] + noise * rng.standard_normal(DIM)) + db.command( + "sql", + "CREATE VERTEX Fact SET fid = ?, tenant = ?, text = ?, embedding = ?", + fid, + tenant, + text, + arcadedb.to_java_float_array(emb.astype(np.float32).tolist()), + ) + for name in (a, b): + mention_edges.append((fid, entity_ids[(tenant, topic, str(name))])) + fid += 1 + for f, e in mention_edges: + db.command( + "sql", + "CREATE EDGE MENTIONS FROM (SELECT FROM Fact WHERE fid = ?) " + "TO (SELECT FROM Entity WHERE eid = ?)", + f, + e, + ) + return centroids + + +def _unit(v: np.ndarray) -> np.ndarray: + return v / np.linalg.norm(v) + + +def query_vector(centroids: dict) -> list: + """The embedding of 'what coffee does the user like': near the coffee topic.""" + rng = np.random.default_rng(SEED + 1) + v = _unit(centroids["coffee"] + 0.05 * rng.standard_normal(DIM)) + return v.astype(np.float32).tolist() + + +def run(db, centroids: dict, form: str) -> dict: + """Run the statements for REQUESTED_TENANT and return everything the checks need.""" + params = { + "qvec": arcadedb.to_java_float_array(query_vector(centroids)), + "tenant": REQUESTED_TENANT, + "keywords": "espresso grinder", + "k": K, + "top": TOP, + } + + def q(sql, **extra): + every = {**params, **extra} + used = {n: every[n] for n in re.findall(r":(\w+)", sql) if n in every} + return db.query("sql", sql, used).to_list() + + # What goes wrong without a pre-filter: the global top-K, then the tenant test. + post = q("SELECT fid FROM (" + neighbors_sql(None) + ") WHERE tenant = :tenant") + pre = q(neighbors_sql(RESTRICT[form])) + hybrid = q(HYBRID_SQL[form]) + # The same hops without the tenant guard, to show what the guard removes. + unguarded = q( + "SELECT out('MENTIONS').out('RELATES').in('MENTIONS').fid AS fids " + "FROM Fact WHERE fid = :fid", + fid=hybrid[0]["fid"], + )[0]["fids"] + return { + "post_filtered": post, + "pre_filtered": pre, + "hybrid": hybrid, + "unguarded_fids": unguarded, + } + + +def tenants_of(db, fids) -> set: + """Tenants of the given fact ids (empty input: empty set).""" + fids = list(fids) + if not fids: + return set() + rows = db.query("sql", "SELECT tenant FROM Fact WHERE fid IN :fids", {"fids": fids}) + return {r["tenant"] for r in rows.to_list()} + + +def main() -> None: + args = parse_args() + db_path = Path(args.db_path) + if db_path.exists(): + shutil.rmtree(db_path) + db_path.parent.mkdir(parents=True, exist_ok=True) + with arcadedb.create_database(str(db_path)) as db: + centroids = build(db) + counts = { + t: db.query("sql", f"SELECT count(*) AS n FROM {t}").to_list()[0]["n"] + for t in ("Fact", "Entity", "MENTIONS", "RELATES") + } + print(f"memory: {counts}, tenants: {len(TENANTS)}") + print(f"engine: arcadedb_embedded {arcadedb.__version__}, form: {args.form}") + out = run(db, centroids, args.form) + + print(f"\nrequested tenant: {REQUESTED_TENANT}, k = {K}") + print( + f"global top-{K}, then WHERE tenant: {len(out['post_filtered'])} rows " + "(the nearest neighbours belong to other tenants)" + ) + print(f"pre-filtered top-{K} ({args.form}): {len(out['pre_filtered'])} rows") + + print(f"\nhybrid result (vector + full-text + graph), top {TOP}:") + for r in out["hybrid"]: + print(f" fid {r['fid']:>3} rrf {r['score']:.4f} {r['fact']}") + print(f" related entities: {r['related_entities']}") + print(f" related facts: {r['related_fids']}") + + # Checks: every row, and every fact reached by a hop, is the tenant's. + shown = [r["fid"] for r in out["hybrid"]] + related = [f for r in out["hybrid"] for f in r["related_fids"] or []] + checks = { + "pre-filtered rows are all the tenant's": tenants_of( + db, [r["fid"] for r in out["pre_filtered"]] + ), + "hybrid rows are all the tenant's": tenants_of(db, shown), + "hop results are all the tenant's": tenants_of(db, related), + } + for label, found in checks.items(): + assert found == {REQUESTED_TENANT}, (label, found) + print(f"check ok: {label} {sorted(found)}") + assert len(out["pre_filtered"]) == K and len(out["hybrid"]) == TOP + leaked = tenants_of(db, out["unguarded_fids"]) - {REQUESTED_TENANT} + print( + f"without the [tenant = :tenant] hop guard the hops also reach: {sorted(leaked)}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/README.md b/examples/README.md index e167306..e1637ea 100644 --- a/examples/README.md +++ b/examples/README.md @@ -391,6 +391,26 @@ half-applied after an interruption --- +### 🧠 [29_hybrid_memory_retrieval.py](./29_hybrid_memory_retrieval.py) +**Agent Memory | Tenant Pre-filter | Vector + Full-Text + Graph | `vector.fuse`** + +Hybrid retrieval for an agent memory in one SQL statement instead of three stores: +- `vector.neighbors` restricted to one tenant BEFORE the top-k (`where`, ArcadeDB 26.11.1, + or `filter` with a RID subquery, 26.10.2 and later), a `SEARCH_INDEX` full-text match on + the same tenant, both fused by reciprocal rank fusion with `vector.fuse`, then graph hops + (`MENTIONS`, `RELATES`) to related entities and facts, kept inside the tenant +- shows why: a global top-k followed by a tenant test returns 0 rows when the nearest + neighbours belong to another tenant (the behaviour of Cypher `db.index.vector.queryNodes`, + which has no filter) +- checks that every result row, and every fact reached by a hop, belongs to the tenant +- builds a small synthetic memory (fixed-seed numpy embeddings); no data download +- flags: `--form` (`where`, default, needs 26.11.1; or `filter`) and `--db-path` + +**Learn:** How to filter a vector search by tenant before the top-k and fuse it with +full-text and graph context in a single query + +--- + ## 💡 Tips - **Run from examples/ directory** - Always execute examples from `examples/` for correct file paths diff --git a/tests/test_example29_hybrid_memory.py b/tests/test_example29_hybrid_memory.py new file mode 100644 index 0000000..66bbc4a --- /dev/null +++ b/tests/test_example29_hybrid_memory.py @@ -0,0 +1,74 @@ +"""Example 29 at test size: the one-statement hybrid retrieval returns only the +requested tenant's facts, where a global top-k followed by a tenant test returns +none, and the tenant guard keeps the graph hops inside the tenant. + +Both pre-filter forms run: ``filter`` (a RID subquery) and ``where`` (ArcadeDB +26.11.1, the version this package is built against).""" + +import importlib.util +from pathlib import Path + +import pytest + +import arcadedb_embedded as arcadedb + +EXAMPLE_PATH = ( + Path(__file__).resolve().parents[1] / "examples" / "29_hybrid_memory_retrieval.py" +) + + +@pytest.fixture(scope="module") +def example29(): + spec = importlib.util.spec_from_file_location("example29", EXAMPLE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def memory(example29, tmp_path_factory): + path = tmp_path_factory.mktemp("hybrid") / "db" + with arcadedb.create_database(str(path)) as db: + centroids = example29.build(db) + yield db, centroids + + +@pytest.mark.parametrize("form", ["filter", "where"]) +def test_hybrid_rows_and_hops_stay_in_the_tenant(example29, memory, form): + db, centroids = memory + out = example29.run(db, centroids, form) + tenant = example29.REQUESTED_TENANT + + assert len(out["pre_filtered"]) == example29.K + assert len(out["hybrid"]) == example29.TOP + scores = [r["score"] for r in out["hybrid"]] + assert scores == sorted(scores, reverse=True) + assert example29.tenants_of(db, [r["fid"] for r in out["pre_filtered"]]) == {tenant} + assert example29.tenants_of(db, [r["fid"] for r in out["hybrid"]]) == {tenant} + related = [f for r in out["hybrid"] for f in r["related_fids"]] + assert related and example29.tenants_of(db, related) == {tenant} + assert all(r["related_entities"] for r in out["hybrid"]) + + +def test_filtering_after_the_top_k_loses_the_rows(example29, memory): + db, centroids = memory + out = example29.run(db, centroids, "filter") + assert len(out["post_filtered"]) < example29.K # the pre-filter returns K + + +def test_the_forms_agree(example29, memory): + db, centroids = memory + by_filter = example29.run(db, centroids, "filter") + by_where = example29.run(db, centroids, "where") + assert [r["fid"] for r in by_filter["hybrid"]] == [ + r["fid"] for r in by_where["hybrid"] + ] + + +def test_the_hop_guard_matters(example29, memory): + """The memory has one stray cross-tenant RELATES edge; without the guard the + hops of the acme facts reach globex facts.""" + db, centroids = memory + out = example29.run(db, centroids, "filter") + leaked = example29.tenants_of(db, out["unguarded_fids"]) + assert leaked - {example29.REQUESTED_TENANT}