Persist full LightRAG state (KV, vectors, doc status) in Memgraph#214
Merged
Conversation
Add Memgraph-backed KV, vector, and doc-status storage backends so the whole LightRAG working state (not just the entity/relationship graph) lives in Memgraph, matching how the graph already persists. - MemgraphKVStorage, MemgraphVectorStorage, MemgraphDocStatusStorage implement the full BaseKVStorage / BaseVectorStorage / DocStatusStorage interfaces, mirroring the connection style (shared async neo4j driver, MEMGRAPH_URI/USERNAME/PASSWORD/DATABASE env vars, config.ini fallback) and error handling of lightrag's built-in graph MemgraphStorage. - Distinct label namespaces (LightRAGKV_*, LightRAGVector_*, LightRAGDocStatus_*) so KV/vector/doc-status nodes never collide with the graph's entity nodes. Vector store uses Memgraph's native cosine vector index + vector_search.search. - register_memgraph_storages() patches lightrag's STORAGES / STORAGE_IMPLEMENTATIONS / STORAGE_ENV_REQUIREMENTS with absolute module paths so verify_storage_implementation accepts the backends by name; idempotent. - Wire MemgraphLightRAGWrapper to route KV + doc-status (always) and vector (when embeddings enabled) to Memgraph by default, gated behind a new full_memgraph_persistence flag; disable_embeddings still falls back to NanoVectorDB for vectors only. - Bump package to 0.2.0, update READMEs, add tests exercising registration, verification and instantiation without a live server. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XeWmYhcoGHac3Q1QMS8tZN
MemgraphDocStatusStorage.get_all_status_counts() returned only per-status counts, unlike every reference DocStatusStorage impl (Json/Redis/Mongo/ OpenSearch), which add an "all" grand-total key. The status-counts API (document_routes) relies on that key. Mirror the reference contract. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XeWmYhcoGHac3Q1QMS8tZN
…x 1.5.x compat Add a session-scoped, live-Memgraph integration test suite exercising the KV, vector and doc-status backends against a real server (one instance reused for the whole suite; per-test unique workspaces + drop() for isolation; skips cleanly when no server is reachable). Uses a deterministic fake embedding function so no external embedding API is needed, and proves the vector path: vector_search.search nearest-neighbour ordering and cosine_better_than_threshold filtering. vector_impl: pass the vector index name to vector_search.search as a string-literal argument (matching the toolbox NodeVectorSearchTool and the unstructured2graph graphrag example) instead of a Cypher parameter; the index name is already sanitized so this is injection safe. Wire the test into CI: the shared tests.yaml lightrag job now waits for the Memgraph MAGE container via a Bolt readiness check and points the backends at it via MEMGRAPH_URI so the integration tests run. Fix lightrag-hku 1.5.x compatibility (1.5.4 added abstract methods to DocStatusStorage): implement get_doc_by_file_basename and get_doc_by_content_hash with indexed Cypher lookups (content_hash promoted to an indexed node property), and update get_docs_paginated for the new status_filters parameter. Cap the dependency at lightrag-hku>=1.5,<1.6. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XeWmYhcoGHac3Q1QMS8tZN
A bare DETACH DELETE removed nodes from the graph but left dangling entries in Memgraph's native (HNSW) vector index. The next vector_search.search then returned a handle to a deleted node and reading a property off it raised 50N42 "Trying to get a property from a deleted object", which the broad except in query() swallowed by returning [] -- wiping the entire result set, including live matches. Two layers of defence: - Index eviction on delete: SET n.embedding = NULL before DETACH DELETE in every deletion path (delete, delete_entity_relation, drop) so the entry leaves the vector index (removing the indexed property removes the vector from the index) rather than dangling. - Resilient query(): re-MATCH live nodes by identity and read properties only off those, never off the raw procedure handle. A dangling candidate matches no live node and is silently excluded; identity comparison touches no properties so it cannot raise 50N42. The broad except no longer returns [] on any error -- it logs and re-raises so real failures are not masked as empty results. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XeWmYhcoGHac3Q1QMS8tZN
antejavor
requested changes
Jul 14, 2026
| "MemgraphKVStorage", | ||
| "MemgraphLightRAGWrapper", | ||
| "MemgraphVectorStorage", | ||
| "register_memgraph_storages", |
Contributor
There was a problem hiding this comment.
This is a strange place to put pural, since memgraph_storage is single.
Comment on lines
+20
to
+65
| from neo4j import AsyncDriver, AsyncGraphDatabase | ||
|
|
||
| # Mirror lightrag.kg.memgraph_impl: allow a config.ini [memgraph] section to | ||
| # provide connection defaults when the environment variables are not set. | ||
| _config = configparser.ConfigParser() | ||
| _config.read("config.ini", "utf-8") | ||
|
|
||
| # One shared driver per event loop, keyed by id(loop). Reusing the driver | ||
| # across storages avoids the overhead of a driver-per-call and keeps a single | ||
| # connection pool, while keying by loop id keeps async test suites (which spin | ||
| # up a fresh loop per test) from reusing a driver bound to a closed loop. | ||
| _drivers: dict[int, AsyncDriver] = {} | ||
| _databases: dict[int, str] = {} | ||
|
|
||
|
|
||
| def read_connection_config() -> tuple[str, str, str, str]: | ||
| """Read Memgraph connection settings (env first, then config.ini fallback). | ||
|
|
||
| Returns: | ||
| Tuple of ``(uri, username, password, database)``. | ||
| """ | ||
| uri = os.environ.get( | ||
| "MEMGRAPH_URI", | ||
| _config.get("memgraph", "uri", fallback="bolt://localhost:7687"), | ||
| ) | ||
| username = os.environ.get("MEMGRAPH_USERNAME", _config.get("memgraph", "username", fallback="")) | ||
| password = os.environ.get("MEMGRAPH_PASSWORD", _config.get("memgraph", "password", fallback="")) | ||
| database = os.environ.get( | ||
| "MEMGRAPH_DATABASE", | ||
| _config.get("memgraph", "database", fallback="memgraph"), | ||
| ) | ||
| return uri, username, password, database | ||
|
|
||
|
|
||
| async def get_driver() -> AsyncDriver: | ||
| """Return the shared async driver for the current event loop, creating it if needed.""" | ||
| loop_id = id(asyncio.get_running_loop()) | ||
| driver = _drivers.get(loop_id) | ||
| if driver is None: | ||
| uri, username, password, database = read_connection_config() | ||
| driver = AsyncGraphDatabase.driver(uri, auth=(username, password)) | ||
| _drivers[loop_id] = driver | ||
| _databases[loop_id] = database | ||
| logger.info(f"Opened shared Memgraph driver for LightRAG storages at {uri}") | ||
| return driver | ||
|
|
Contributor
There was a problem hiding this comment.
Can this be part of the https://github.com/memgraph/ai-toolkit/blob/main/memgraph-toolbox/src/memgraph_toolbox/api/memgraph.py?
Why we have a specific way of getting a driver here? I think this should depend on the toolbox, anything in the unstrucured2graph should depend on it.
Comment on lines
60
to
62
| if self.disable_embeddings: | ||
| lightrag_kwargs["embedding_func"] = _dummy_embedding_func(dim=1) | ||
| lightrag_kwargs["vector_storage"] = "NanoVectorDBStorage" |
Contributor
There was a problem hiding this comment.
Does this makes sense, I don't want to have nanovector since I have the vector storage in Memgraph.
| # doc-status always go to Memgraph; the vector store only goes to | ||
| # Memgraph when embeddings are enabled (a real vector index needs a | ||
| # real embedding dimension), otherwise it keeps the local | ||
| # NanoVectorDB fallback set above. Explicit caller overrides win. |
Contributor
There was a problem hiding this comment.
Don't want the fallback.
Comment on lines
+284
to
+285
| driver = await get_driver() | ||
| async with driver.session(database=get_database(), default_access_mode="READ") as session: |
Contributor
There was a problem hiding this comment.
This is also driver, should come from the toolbox.
…ctors, rename - Add additive AsyncMemgraph client to memgraph-toolbox (neo4j AsyncGraphDatabase) mirroring the sync Memgraph shape; sync class and memgraph_env untouched. Add a construction/config unit test that skips the live round-trip cleanly when no server is reachable. - Source the async driver in lightrag-memgraph's _connection.py from AsyncMemgraph instead of a bespoke AsyncGraphDatabase driver, keeping per-event-loop pooling and the same get_driver/get_database/close_driver signatures. Reconcile env-var names by reading both MEMGRAPH_URI/URL and MEMGRAPH_USERNAME/USER (preferring the LightRAG names to match the graph backend) and passing resolved values explicitly to AsyncMemgraph. - Drop the NanoVectorDB fallback and the disable_embeddings mode: vectors always persist to Memgraph's native vector index (a real embedding_func is required). Update core.py, README, example.py, unstructured2graph example and tests accordingly. - Rename register_memgraph_storages -> register_memgraph_storage everywhere. - CI: set both MEMGRAPH_URL/USER and MEMGRAPH_URI/USERNAME in the test-lightrag-memgraph job so the graph backend and toolbox-backed storages resolve to the same instance. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XeWmYhcoGHac3Q1QMS8tZN
antejavor
requested changes
Jul 15, 2026
Comment on lines
+214
to
+215
| MEMGRAPH_URI: bolt://localhost:7687 | ||
| MEMGRAPH_URL: bolt://localhost:7687 |
Contributor
There was a problem hiding this comment.
Looks like this is the reason -> https://github.com/memgraph/ai-toolkit/pull/214/changes#r3585308930
Comment on lines
+11
to
+13
| ``MEMGRAPH_URI`` / ``MEMGRAPH_USERNAME`` / ``MEMGRAPH_PASSWORD`` / | ||
| ``MEMGRAPH_DATABASE``, while the toolbox reads | ||
| ``MEMGRAPH_URL`` / ``MEMGRAPH_USER`` / ``MEMGRAPH_PASSWORD`` / |
Contributor
There was a problem hiding this comment.
Ok so the URI and URL is different because of MemgraphSTorage that is part of the LightRAG lib.
Comment on lines
+33
to
+34
| if TYPE_CHECKING: | ||
| from neo4j import AsyncDriver |
Contributor
There was a problem hiding this comment.
Should this happen in toolbox?
Comment on lines
+1
to
+146
| """Shared Memgraph connection helpers for the LightRAG storage backends. | ||
|
|
||
| All three storage backends (KV, vector, doc-status) reuse a single async Bolt | ||
| driver per running event loop instead of opening one driver per storage | ||
| instance or per call. The driver comes from | ||
| ``memgraph_toolbox.api.memgraph.AsyncMemgraph`` so this integration shares the | ||
| toolbox's connection contract instead of maintaining a bespoke driver. | ||
|
|
||
| Connection settings mirror LightRAG's own graph backend | ||
| (``lightrag.kg.memgraph_impl.MemgraphStorage``), which reads | ||
| ``MEMGRAPH_URI`` / ``MEMGRAPH_USERNAME`` / ``MEMGRAPH_PASSWORD`` / | ||
| ``MEMGRAPH_DATABASE``, while the toolbox reads | ||
| ``MEMGRAPH_URL`` / ``MEMGRAPH_USER`` / ``MEMGRAPH_PASSWORD`` / | ||
| ``MEMGRAPH_DATABASE``. To avoid a split-brain where the graph backend and these | ||
| storages point at different instances, both name variants are read here | ||
| (preferring the LightRAG names so we match the graph backend) and the resolved | ||
| values are passed EXPLICITLY to ``AsyncMemgraph`` so env-name drift cannot cause | ||
| a mismatch. The ``config.ini [memgraph]`` fallback is preserved. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import asyncio | ||
| import configparser | ||
| import os | ||
| import re | ||
| from typing import TYPE_CHECKING | ||
|
|
||
| from lightrag.utils import logger | ||
|
|
||
| from memgraph_toolbox.api.memgraph import AsyncMemgraph | ||
|
|
||
| if TYPE_CHECKING: | ||
| from neo4j import AsyncDriver | ||
|
|
||
| # Mirror lightrag.kg.memgraph_impl: allow a config.ini [memgraph] section to | ||
| # provide connection defaults when the environment variables are not set. | ||
| _config = configparser.ConfigParser() | ||
| _config.read("config.ini", "utf-8") | ||
|
|
||
| # One shared AsyncMemgraph client per event loop, keyed by id(loop). Reusing the | ||
| # driver across storages avoids the overhead of a driver-per-call and keeps a | ||
| # single connection pool, while keying by loop id keeps async test suites (which | ||
| # spin up a fresh loop per test) from reusing a driver bound to a closed loop. | ||
| _clients: dict[int, AsyncMemgraph] = {} | ||
|
|
||
|
|
||
| def _first_env(*names: str) -> str | None: | ||
| """Return the first environment variable that is set (even to ""), else None.""" | ||
| for name in names: | ||
| value = os.environ.get(name) | ||
| if value is not None: | ||
| return value | ||
| return None | ||
|
|
||
|
|
||
| def read_connection_config() -> tuple[str, str, str, str]: | ||
| """Read Memgraph connection settings (env first, then config.ini fallback). | ||
|
|
||
| Both the LightRAG env-var names (``MEMGRAPH_URI`` / ``MEMGRAPH_USERNAME``) | ||
| and the toolbox names (``MEMGRAPH_URL`` / ``MEMGRAPH_USER``) are honoured, | ||
| preferring the LightRAG names so these storages resolve to the same instance | ||
| as LightRAG's built-in graph backend. | ||
|
|
||
| Returns: | ||
| Tuple of ``(uri, username, password, database)``. | ||
| """ | ||
| uri = _first_env("MEMGRAPH_URI", "MEMGRAPH_URL") | ||
| if uri is None: | ||
| uri = _config.get("memgraph", "uri", fallback="bolt://localhost:7687") | ||
|
|
||
| username = _first_env("MEMGRAPH_USERNAME", "MEMGRAPH_USER") | ||
| if username is None: | ||
| username = _config.get("memgraph", "username", fallback="") | ||
|
|
||
| password = _first_env("MEMGRAPH_PASSWORD") | ||
| if password is None: | ||
| password = _config.get("memgraph", "password", fallback="") | ||
|
|
||
| database = _first_env("MEMGRAPH_DATABASE") | ||
| if database is None: | ||
| database = _config.get("memgraph", "database", fallback="memgraph") | ||
|
|
||
| return uri, username, password, database | ||
|
|
||
|
|
||
| def _get_client() -> AsyncMemgraph: | ||
| """Return the shared AsyncMemgraph client for the current event loop, creating it if needed.""" | ||
| loop_id = id(asyncio.get_running_loop()) | ||
| client = _clients.get(loop_id) | ||
| if client is None: | ||
| uri, username, password, database = read_connection_config() | ||
| # Pass values EXPLICITLY so AsyncMemgraph does not re-resolve from the | ||
| # toolbox env-var names and risk a mismatch with the graph backend. | ||
| client = AsyncMemgraph( | ||
| url=uri, | ||
| username=username, | ||
| password=password, | ||
| database=database, | ||
| user_agent="lightrag-memgraph", | ||
| ) | ||
| _clients[loop_id] = client | ||
| logger.info(f"Opened shared Memgraph driver for LightRAG storages at {uri}") | ||
| return client | ||
|
|
||
|
|
||
| async def get_driver() -> AsyncDriver: | ||
| """Return the shared async driver for the current event loop, creating it if needed.""" | ||
| return _get_client().driver | ||
|
|
||
|
|
||
| def get_database() -> str: | ||
| """Return the database name bound to the current event loop's driver.""" | ||
| loop_id = id(asyncio.get_running_loop()) | ||
| client = _clients.get(loop_id) | ||
| if client is not None: | ||
| return client.database | ||
| return read_connection_config()[3] | ||
|
|
||
|
|
||
| async def close_driver() -> None: | ||
| """Close the shared driver bound to the current event loop, if any.""" | ||
| loop_id = id(asyncio.get_running_loop()) | ||
| client = _clients.pop(loop_id, None) | ||
| if client is not None: | ||
| await client.close() | ||
| logger.info("Closed shared Memgraph driver for LightRAG storages") | ||
|
|
||
|
|
||
| def sanitize_label(value: str) -> str: | ||
| """Escape a value for safe use as a backtick-quoted Cypher label identifier. | ||
|
|
||
| Backticks are doubled to prevent Cypher injection; all other characters are | ||
| preserved. The result is intended to be wrapped in backticks, e.g. | ||
| ``MATCH (n:`{label}`)``. Mirrors ``MemgraphStorage._get_workspace_label``. | ||
| """ | ||
| return value.strip().replace("`", "``") | ||
|
|
||
|
|
||
| def sanitize_index_name(value: str) -> str: | ||
| """Coerce a value into a bare (unquoted) identifier for a vector index name. | ||
|
|
||
| Memgraph vector-index names are unquoted identifiers, so every character | ||
| outside ``[A-Za-z0-9_]`` is replaced with an underscore. | ||
| """ | ||
| return re.sub(r"[^A-Za-z0-9_]", "_", value) |
Contributor
There was a problem hiding this comment.
Think about this connection code, what needs to end up in the toolbox, and what should be here in light rag.
|
|
||
| def __post_init__(self): | ||
| self._validate_embedding_func() | ||
| workspace = os.environ.get("MEMGRAPH_WORKSPACE") or self.workspace or "base" |
Contributor
There was a problem hiding this comment.
From where is this coming the workspace env?
| embedding = np.asarray(embedded[0], dtype=float).tolist() | ||
|
|
||
| driver = await get_driver() | ||
| async with driver.session(database=get_database(), default_access_mode="READ") as session: |
Contributor
There was a problem hiding this comment.
Read the Memgraph docs on vector search -> https://memgraph.com/docs/querying/vector-search and see those the live makes sense.
This could be a problem of the READ_UNCOMMITTED
| import traceback | ||
|
|
||
| from lightrag.llm.anthropic import anthropic_complete | ||
| from lightrag.llm.openai import openai_embed |
| max_parallel_insert=8, | ||
| llm_model_func=anthropic_complete, | ||
| llm_model_name="claude-haiku-4-5", | ||
| embedding_func=openai_embed, |
Contributor
There was a problem hiding this comment.
Function is from openai and model from Antropic.
|
|
||
| dependencies = [ | ||
| "lightrag-hku>=1.4.14", | ||
| "lightrag-hku>=1.5,<1.6", |
Contributor
There was a problem hiding this comment.
Let's remove the upper bound.
| working_dir="./lightrag_storage", | ||
| llm_model_func=anthropic_complete, | ||
| llm_model_name="claude-3-5-sonnet-20241022", # or claude-3-haiku-20240307, etc. | ||
| embedding_func=openai_embed, # Anthropic has no embeddings; supply one |
Contributor
There was a problem hiding this comment.
Is there a way to make this work -> https://memgraph.com/docs/advanced-algorithms/available-algorithms/embeddings?
A. Toolbox owns connection config + single env set: - _connection.py: drop bespoke env reading (read_connection_config, _first_env) and the config.ini fallback; construct AsyncMemgraph() with no connection args so the toolbox's memgraph_env is the single source of truth. Keep the per-event-loop driver cache, get_driver/get_database/ close_driver, and the sanitizers. - core.py: move the env-name bridge into initialize() and add the MEMGRAPH_USER -> MEMGRAPH_USERNAME bridge alongside URL -> URI, without clobbering explicit LightRAG names. - registry.py: STORAGE_ENV_REQUIREMENTS now requires MEMGRAPH_URL (canonical toolbox name). - tests.yaml: collapse the double env set to the canonical toolbox names only (MEMGRAPH_URL/USER/PASSWORD/DATABASE); core.py bridges to LightRAG. - Integration test resolves connection via the toolbox memgraph_env. B. pyproject.toml: drop the lightrag-hku <1.6 upper bound. C. example.py: make the example OpenAI-only (gpt_4o_mini_complete + openai_embed) for a self-consistent minimal example; comment documents the Claude + separate-embedder alternative. D. vector_impl.py: tighten the vector-search comment to state the native vector index runs at READ_UNCOMMITTED with deferred GC (the source of 50N42) and that default_access_mode="READ" is only a routing hint, not an isolation control. No Cypher/logic change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XeWmYhcoGHac3Q1QMS8tZN
antejavor
marked this pull request as ready for review
July 16, 2026 07:59
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Requested by Ante Javor · Slack thread
What changes
Before: running
unstructured2graph, only the knowledge graph landed in Memgraph. LightRAG's key-value store, its embeddings/vectors, and its document-processing status were written to local files in the working dir (JsonKVStorage,NanoVectorDBStorage,JsonDocStatusStorage), so that state wasn't durable in Memgraph and couldn't be queried there.After: LightRAG's full working state persists in Memgraph by default — key-value data, vectors, and doc status, alongside the graph. Nothing is left on local disk unless you opt out.
How
Adds three Memgraph-backed LightRAG storage backends to the vendored
lightrag-memgraphpackage (integrations/lightrag-memgraph/), a registry patch so LightRAG accepts them by name, and wrapper wiring so they're used by default. Written from scratch, no dependency on any external patch.MemgraphKVStorage(BaseKVStorage)— nodes labelledLightRAGKV_<workspace>_<namespace>, value dict stored as a JSONdataproperty, label index onid.MemgraphVectorStorage(BaseVectorStorage)— nodes labelledLightRAGVector_<workspace>_<namespace>; embeddings on anembeddingproperty backed by a native Memgraph vector index (metric: cos, dimension from the embedding func);query()usesCALL vector_search.search(...)filtered bycosine_better_than_threshold.MemgraphDocStatusStorage(DocStatusStorage)— nodes labelledLightRAGDocStatus_<workspace>;status/track_id/file_path/created_at/updated_atpromoted to indexed props for filter/sort.initialize().register_memgraph_storages()): idempotently patches LightRAG'sSTORAGES,STORAGE_IMPLEMENTATIONS, andSTORAGE_ENV_REQUIREMENTSwith absolute module paths (LightRAG imports impls withpackage="lightrag", so relative paths wouldn't resolve).core.py): the wrapper registers the backends beforeLightRAG(...)and setskv_storage/doc_status_storage/vector_storageviasetdefault(explicit caller overrides still win). New flagfull_memgraph_persistence: bool = Trueopts back out to graph-only. Withdisable_embeddings=True, KV + doc status still go to Memgraph and vectors fall back to NanoVectorDB.MEMGRAPH_URI/MEMGRAPH_USERNAME/MEMGRAPH_PASSWORD/MEMGRAPH_DATABASE/MEMGRAPH_WORKSPACE(and the existingMEMGRAPH_URL→MEMGRAPH_URIbridge), with the sameconfig.ini [memgraph]fallback.Testing
12 unit/sanity tests pass (registration idempotency + absolute paths,
verify_storage_implementationaccepts all three names, every class instantiates with no missing abstract methods, wiring, doc-status round-trip and status-count totals).ruffclean.A live-Memgraph integration test wired into CI (reusing one Memgraph instance across the suite) is being added in a follow-up commit on this branch — it exercises the vector-index DDL and
vector_search.searchround-trip against a real server, which can't be validated in a unit-only env.Generated by Claude Code