lightrag-memgraph 0.3.0__tar.gz → 0.3.2__tar.gz

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: lightrag-memgraph
-Version: 0.3.0
+Version: 0.3.2
 Summary: LightRAG integration with Memgraph
 Project-URL: Homepage, https://github.com/memgraph/ai-toolkit
 Project-URL: Source, https://github.com/memgraph/ai-toolkit/tree/main/integrations/lightrag-memgraph
@@ -18,7 +18,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.10
 Requires-Dist: anthropic>=0.40.0
-Requires-Dist: lightrag-hku<1.6,>=1.5
+Requires-Dist: lightrag-hku==1.5.4
 Requires-Dist: memgraph-toolbox>=0.1.11
 Requires-Dist: numpy>=1.22.0
 Requires-Dist: openai>=1.0.0
@@ -26,7 +26,6 @@ Requires-Dist: voyageai>=0.2.0
 Provides-Extra: dev
 Requires-Dist: black>=26.3.1; extra == 'dev'
 Requires-Dist: isort>=5.0.0; extra == 'dev'
-Requires-Dist: mypy>=1.0.0; extra == 'dev'
 Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
 Requires-Dist: pytest>=9.0.3; extra == 'dev'
 Description-Content-Type: text/markdown
@@ -85,6 +84,12 @@ asyncio.run(main())
 
 See `example.py` in this repo for a full run with sample texts and graph output.
 
+`initialize(**lightrag_kwargs)` forwards anything to `LightRAG(...)`, and explicit
+values override the Memgraph-storage defaults — e.g. `max_parallel_insert=8`, or
+`addon_params={"entity_types_guidance": "..."}` to steer entity extraction toward a
+specific vocabulary. [unstructured2graph](../../unstructured2graph/README.md#entity-typing--ontology)
+produces that `addon_params` from an ontology file via `Ontology.addon_params()`.
+
 ## Storage / persistence
 
 By default the wrapper persists **LightRAG's entire working state** into
@@ -184,7 +189,7 @@ https://platform.claude.com/docs/en/about-claude/models.
    await wrapper.initialize(
        working_dir="./lightrag_storage",
        llm_model_func=anthropic_complete,
-       llm_model_name="claude-3-5-sonnet-20241022",  # or claude-3-haiku-20240307, etc.
+       llm_model_name="claude-sonnet-4-20250514",  # any current Claude model id
        # embedding_func omitted: defaults to Memgraph's local sentence-transformer,
        # see Embeddings above. Anthropic itself has no embeddings API.
    )

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/README.md

@@ -52,6 +52,12 @@ asyncio.run(main())
 
 See `example.py` in this repo for a full run with sample texts and graph output.
 
+`initialize(**lightrag_kwargs)` forwards anything to `LightRAG(...)`, and explicit
+values override the Memgraph-storage defaults — e.g. `max_parallel_insert=8`, or
+`addon_params={"entity_types_guidance": "..."}` to steer entity extraction toward a
+specific vocabulary. [unstructured2graph](../../unstructured2graph/README.md#entity-typing--ontology)
+produces that `addon_params` from an ontology file via `Ontology.addon_params()`.
+
 ## Storage / persistence
 
 By default the wrapper persists **LightRAG's entire working state** into
@@ -151,7 +157,7 @@ https://platform.claude.com/docs/en/about-claude/models.
    await wrapper.initialize(
        working_dir="./lightrag_storage",
        llm_model_func=anthropic_complete,
-       llm_model_name="claude-3-5-sonnet-20241022",  # or claude-3-haiku-20240307, etc.
+       llm_model_name="claude-sonnet-4-20250514",  # any current Claude model id
        # embedding_func omitted: defaults to Memgraph's local sentence-transformer,
        # see Embeddings above. Anthropic itself has no embeddings API.
    )

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "lightrag-memgraph"
-version = "0.3.0"
+version = "0.3.2"
 description = "LightRAG integration with Memgraph"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -22,10 +22,15 @@ classifiers = [
 ]
 
 dependencies = [
-    # Capped below 1.6: 1.5.4 broke this integration's DocStatusStorage once
-    # already (see registry.py's compat guard). Bump only after verifying
-    # compatibility with the new release.
-    "lightrag-hku>=1.5,<1.6",
+    # Pinned to exactly 1.5.4: 1.5.5 (2026-07-31) added new @abstractmethod
+    # requirements to DocStatusStorage (a "memory-bounded scheduling" API --
+    # get_docs_by_statuses_page/get_docs_by_ids/get_full_docs_by_ids/
+    # resolve_doc_source_strict) that MemgraphDocStatusStorage doesn't
+    # implement yet, so LightRAG() construction raises TypeError. See #250.
+    # Also still capped below 1.6 for the separate reason registry.py's
+    # compat guard documents (lightrag.kg internals this integration
+    # monkey-patches). Bump only after implementing the new interface.
+    "lightrag-hku==1.5.4",
     "memgraph-toolbox>=0.1.11",
     "numpy>=1.22.0",
     "anthropic>=0.40.0",
@@ -39,7 +44,6 @@ dev = [
     "pytest-asyncio>=0.21.0",
     "black>=26.3.1",
     "isort>=5.0.0",
-    "mypy>=1.0.0",
 ]
 
 [project.urls]
@@ -66,9 +70,3 @@ target-version = ['py310']
 [tool.isort]
 profile = "black"
 line_length = 88
-
-[tool.mypy]
-python_version = "3.10"
-warn_return_any = true
-warn_unused_configs = true
-disallow_untyped_defs = true

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/src/lightrag_memgraph/__init__.py

@@ -25,7 +25,7 @@ def _patch_anthropic() -> None:
         import os
         from collections.abc import AsyncIterator
         from importlib.metadata import version as get_version
-        from typing import Any
+        from typing import Any, cast
 
         import lightrag.llm.anthropic as _mod
         from anthropic import (
@@ -121,8 +121,14 @@ def _patch_anthropic() -> None:
                 parts.append(content)
             return "".join(parts)
 
-        _wrapped._lightrag_memgraph_patched = True  # type: ignore[attr-defined]
-        _mod.anthropic_complete_if_cache = _wrapped
+        # _wrapped is a plain function object; cast(..., x) is the identity function at runtime,
+        # so this sets/reads the real attribute -- it's just how we spell "let me stash a marker
+        # attribute on a function" past a type checker that doesn't model dynamic attributes on
+        # function objects. Likewise, _wrapped's signature (no image_inputs=) is intentionally
+        # narrower than anthropic_complete_if_cache's -- it's real dynamic monkey-patching of a
+        # third-party module attribute, not something a structural type can express.
+        cast("Any", _wrapped)._lightrag_memgraph_patched = True
+        _mod.anthropic_complete_if_cache = cast("Any", _wrapped)
     except Exception:
         pass
 

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/src/lightrag_memgraph/_connection.py

@@ -13,7 +13,7 @@ from __future__ import annotations
 import asyncio
 import os
 import re
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from lightrag.utils import logger
 
@@ -21,6 +21,7 @@ from memgraph_toolbox.api.memgraph import AsyncMemgraph, memgraph_env
 
 if TYPE_CHECKING:
     from neo4j import AsyncDriver, AsyncSession
+    from typing_extensions import LiteralString
 
 # Keyed by event-loop id so tests spinning up a fresh loop per test don't
 # reuse a driver bound to a closed loop.
@@ -71,6 +72,19 @@ def sanitize_index_name(value: str) -> str:
     return re.sub(r"[^A-Za-z0-9_]", "_", value)
 
 
+def as_literal_query(text: str) -> LiteralString:
+    """Assert that a dynamically-built Cypher string is safe to pass to ``AsyncSession.run``.
+
+    neo4j types ``run(query=...)`` as ``LiteralString | Query`` to steer callers away from
+    interpolating untrusted input into query text -- ``LiteralString`` can only be produced
+    from literal syntax, so the type checker can't see that a runtime f-string is safe. Every
+    query built in this package only interpolates our own backtick-escaped identifiers
+    (``sanitize_label``/``sanitize_index_name``) or passes values as bound parameters (``$id``
+    etc.), never raw external input concatenated into query text, so the cast is safe here.
+    """
+    return cast("LiteralString", text)
+
+
 def resolve_workspace(workspace: str | None) -> str:
     """Resolve the effective workspace for a storage instance.
 
@@ -84,7 +98,7 @@ def resolve_workspace(workspace: str | None) -> str:
 async def create_index(session: AsyncSession, *, label: str, prop: str, workspace: str) -> None:
     """CREATE INDEX on :`label`(prop), swallowing "already exists" (Memgraph has no IF NOT EXISTS)."""
     try:
-        await (await session.run(f"CREATE INDEX ON :`{label}`({prop})")).consume()
+        await (await session.run(as_literal_query(f"CREATE INDEX ON :`{label}`({prop})"))).consume()
     except Exception as e:  # index may already exist, which is not an error
         logger.warning(f"[{workspace}] Index creation on :`{label}`({prop}) may have failed or already exists: {e}")
 
@@ -106,7 +120,7 @@ class MemgraphCrudMixin:
     async def is_empty(self) -> bool:
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
-            result = await session.run(f"MATCH (n:`{self._label}`) RETURN n LIMIT 1")
+            result = await session.run(as_literal_query(f"MATCH (n:`{self._label}`) RETURN n LIMIT 1"))
             record = await result.single()
             await result.consume()
             return record is None
@@ -118,13 +132,15 @@ class MemgraphCrudMixin:
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"""
-                UNWIND $keys AS k
-                OPTIONAL MATCH (n:`{self._label}` {{id: k}})
-                WITH k, n
-                WHERE n IS NULL
-                RETURN k
-                """,
+                as_literal_query(
+                    f"""
+                    UNWIND $keys AS k
+                    OPTIONAL MATCH (n:`{self._label}` {{id: k}})
+                    WITH k, n
+                    WHERE n IS NULL
+                    RETURN k
+                    """
+                ),
                 keys=list(keys),
             )
             missing = {record["k"] async for record in result}
@@ -138,11 +154,13 @@ class MemgraphCrudMixin:
         async with driver.session(database=get_database()) as session:
             await (
                 await session.run(
-                    f"""
-                    UNWIND $ids AS target_id
-                    MATCH (n:`{self._label}` {{id: target_id}})
-                    DETACH DELETE n
-                    """,
+                    as_literal_query(
+                        f"""
+                        UNWIND $ids AS target_id
+                        MATCH (n:`{self._label}` {{id: target_id}})
+                        DETACH DELETE n
+                        """
+                    ),
                     ids=list(ids),
                 )
             ).consume()
@@ -151,7 +169,7 @@ class MemgraphCrudMixin:
         try:
             driver = await get_driver()
             async with driver.session(database=get_database()) as session:
-                await (await session.run(f"MATCH (n:`{self._label}`) DETACH DELETE n")).consume()
+                await (await session.run(as_literal_query(f"MATCH (n:`{self._label}`) DETACH DELETE n"))).consume()
             logger.info(f"[{self.workspace}] Dropped Memgraph storage for {self.namespace}")
             return {"status": "success", "message": "data dropped"}
         except Exception as e:

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/src/lightrag_memgraph/core.py

@@ -2,7 +2,7 @@ import logging
 import os
 
 from lightrag import LightRAG
-from lightrag.kg.shared_storage import initialize_pipeline_status
+from lightrag.kg.shared_storage import finalize_share_data, initialize_pipeline_status
 from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
 from lightrag.utils import logger, setup_logger
 
@@ -101,6 +101,20 @@ class MemgraphLightRAGWrapper:
             raise RuntimeError("LightRAG not initialized. Call initialize() first.")
         return self.rag
 
+    @property
+    def workspace(self) -> str:
+        """The Memgraph label LightRAG writes every extracted entity under.
+
+        Exposed directly (rather than making callers reach through
+        get_lightrag().chunk_entity_relation_graph) since it's the one piece
+        of LightRAG's internal state consumers outside this package actually
+        need -- e.g. unstructured2graph's LightRAGBackend.workspace_label.
+
+        Raises:
+            RuntimeError: if initialize() hasn't been called yet.
+        """
+        return self.get_lightrag().chunk_entity_relation_graph.workspace
+
     # https://github.com/HKUDS/LightRAG/blob/main/lightrag/lightrag.py
     async def ainsert(self, **kwargs) -> None:
         """
@@ -114,6 +128,22 @@ class MemgraphLightRAGWrapper:
         await self.rag.ainsert(**kwargs)
 
     async def afinalize(self) -> None:
+        """
+        Finalize storages and reset LightRAG's process-global shared-storage
+        state (locks, shared dicts). Without the latter reset, a second
+        LightRAG instance created later in the same process reuses asyncio
+        locks bound to this instance's (by-then-closed) event loop and raises
+        "bound to a different event loop" -- e.g. two real-LightRAG tests
+        run back-to-back under pytest-asyncio's function-scoped event loop.
+
+        The reset runs in a finally block: if finalize_storages() raises
+        (e.g. a transient network error), the shared-storage state would
+        otherwise stay registered as initialized and break every subsequent
+        LightRAG instance in this process, not just this one's cleanup.
+        """
         if self.rag is None:
             raise RuntimeError("LightRAG not initialized. Call initialize() first.")
-        await self.rag.finalize_storages()
+        try:
+            await self.rag.finalize_storages()
+        finally:
+            finalize_share_data()

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/src/lightrag_memgraph/docstatus_impl.py

@@ -11,13 +11,14 @@ from __future__ import annotations
 import asyncio
 import json
 from dataclasses import dataclass
-from typing import Any, final
+from typing import Any, cast, final
 
 from lightrag.base import DocProcessingStatus, DocStatus, DocStatusStorage
 from lightrag.utils import logger
 
 from ._connection import (
     MemgraphCrudMixin,
+    as_literal_query,
     close_driver,
     create_index,
     get_database,
@@ -74,7 +75,7 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"MATCH (n:`{self._label}` {{id: $id}}) RETURN n.data AS data",
+                as_literal_query(f"MATCH (n:`{self._label}` {{id: $id}}) RETURN n.data AS data"),
                 id=id,
             )
             record = await result.single()
@@ -90,11 +91,13 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"""
-                UNWIND $ids AS target_id
-                OPTIONAL MATCH (n:`{self._label}` {{id: target_id}})
-                RETURN target_id AS id, n.data AS data
-                """,
+                as_literal_query(
+                    f"""
+                    UNWIND $ids AS target_id
+                    OPTIONAL MATCH (n:`{self._label}` {{id: target_id}})
+                    RETURN target_id AS id, n.data AS data
+                    """
+                ),
                 ids=list(ids),
             )
             found: dict[str, str] = {}
@@ -102,7 +105,10 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
                 if record["data"] is not None:
                     found[record["id"]] = record["data"]
             await result.consume()
-        return [json.loads(found[i]) if i in found else None for i in ids]
+        # LightRAG's own BaseKVStorage.get_by_ids declares -> list[dict[str, Any]], but its
+        # reference JsonKVStorage impl (and this one, to match) legitimately puts None in the
+        # slot of any id that does not exist -- the upstream annotation itself is imprecise.
+        return cast("list[dict[str, Any]]", [json.loads(found[i]) if i in found else None for i in ids])
 
     async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
         if not data:
@@ -123,11 +129,13 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         async with driver.session(database=get_database()) as session:
             await (
                 await session.run(
-                    f"""
-                    UNWIND $entries AS e
-                    MERGE (n:`{self._label}` {{id: e.id}})
-                    SET n.data = e.data, n += e.indexed
-                    """,
+                    as_literal_query(
+                        f"""
+                        UNWIND $entries AS e
+                        MERGE (n:`{self._label}` {{id: e.id}})
+                        SET n.data = e.data, n += e.indexed
+                        """
+                    ),
                     entries=entries,
                 )
             ).consume()
@@ -141,7 +149,9 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         counts = {status.value: 0 for status in DocStatus}
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
-            result = await session.run(f"MATCH (n:`{self._label}`) RETURN n.status AS status, count(n) AS cnt")
+            result = await session.run(
+                as_literal_query(f"MATCH (n:`{self._label}`) RETURN n.status AS status, count(n) AS cnt")
+            )
             async for record in result:
                 if record["status"] in counts:
                     counts[record["status"]] = record["cnt"]
@@ -164,11 +174,13 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"""
-                MATCH (n:`{self._label}`)
-                WHERE n.status IN $statuses
-                RETURN n.id AS id, n.data AS data
-                """,
+                as_literal_query(
+                    f"""
+                    MATCH (n:`{self._label}`)
+                    WHERE n.status IN $statuses
+                    RETURN n.id AS id, n.data AS data
+                    """
+                ),
                 statuses=status_values,
             )
             records = [record async for record in result]
@@ -179,10 +191,12 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"""
-                MATCH (n:`{self._label}` {{track_id: $track_id}})
-                RETURN n.id AS id, n.data AS data
-                """,
+                as_literal_query(
+                    f"""
+                    MATCH (n:`{self._label}` {{track_id: $track_id}})
+                    RETURN n.id AS id, n.data AS data
+                    """
+                ),
                 track_id=track_id,
             )
             records = [record async for record in result]
@@ -237,7 +251,7 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         async def _count() -> int:
             async with driver.session(database=get_database(), default_access_mode="READ") as session:
                 result = await session.run(
-                    f"MATCH (n:`{self._label}`) {where} RETURN count(n) AS total",
+                    as_literal_query(f"MATCH (n:`{self._label}`) {where} RETURN count(n) AS total"),
                     **filter_params,
                 )
                 record = await result.single()
@@ -247,12 +261,14 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         async def _page() -> list:
             async with driver.session(database=get_database(), default_access_mode="READ") as session:
                 result = await session.run(
-                    f"""
-                    MATCH (n:`{self._label}`) {where}
-                    RETURN n.id AS id, n.data AS data
-                    ORDER BY {sort_expr} {order}
-                    SKIP $skip LIMIT $limit
-                    """,
+                    as_literal_query(
+                        f"""
+                        MATCH (n:`{self._label}`) {where}
+                        RETURN n.id AS id, n.data AS data
+                        ORDER BY {sort_expr} {order}
+                        SKIP $skip LIMIT $limit
+                        """
+                    ),
                     **params,
                 )
                 records = [record async for record in result]
@@ -277,7 +293,7 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"MATCH (n:`{self._label}` {{file_path: $file_path}}) RETURN n.data AS data LIMIT 1",
+                as_literal_query(f"MATCH (n:`{self._label}` {{file_path: $file_path}}) RETURN n.data AS data LIMIT 1"),
                 file_path=file_path,
             )
             record = await result.single()
@@ -293,7 +309,9 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"MATCH (n:`{self._label}` {{file_path: $basename}}) RETURN n.id AS id, n.data AS data LIMIT 1",
+                as_literal_query(
+                    f"MATCH (n:`{self._label}` {{file_path: $basename}}) RETURN n.id AS id, n.data AS data LIMIT 1"
+                ),
                 basename=basename,
             )
             record = await result.single()
@@ -309,7 +327,9 @@ class MemgraphDocStatusStorage(MemgraphCrudMixin, DocStatusStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"MATCH (n:`{self._label}` {{content_hash: $content_hash}}) RETURN n.id AS id, n.data AS data LIMIT 1",
+                as_literal_query(
+                    f"MATCH (n:`{self._label}` {{content_hash: $content_hash}}) RETURN n.id AS id, n.data AS data LIMIT 1"
+                ),
                 content_hash=content_hash,
             )
             record = await result.single()

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/src/lightrag_memgraph/kv_impl.py

@@ -11,13 +11,14 @@ from __future__ import annotations
 import json
 import time
 from dataclasses import dataclass
-from typing import Any, final
+from typing import Any, cast, final
 
 from lightrag.base import BaseKVStorage
 from lightrag.utils import logger
 
 from ._connection import (
     MemgraphCrudMixin,
+    as_literal_query,
     close_driver,
     create_index,
     get_database,
@@ -66,7 +67,7 @@ class MemgraphKVStorage(MemgraphCrudMixin, BaseKVStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"MATCH (n:`{self._label}` {{id: $id}}) RETURN n.data AS data",
+                as_literal_query(f"MATCH (n:`{self._label}` {{id: $id}}) RETURN n.data AS data"),
                 id=id,
             )
             record = await result.single()
@@ -80,11 +81,13 @@ class MemgraphKVStorage(MemgraphCrudMixin, BaseKVStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"""
-                UNWIND $ids AS target_id
-                OPTIONAL MATCH (n:`{self._label}` {{id: target_id}})
-                RETURN target_id AS id, n.data AS data
-                """,
+                as_literal_query(
+                    f"""
+                    UNWIND $ids AS target_id
+                    OPTIONAL MATCH (n:`{self._label}` {{id: target_id}})
+                    RETURN target_id AS id, n.data AS data
+                    """
+                ),
                 ids=list(ids),
             )
             found: dict[str, str] = {}
@@ -92,7 +95,10 @@ class MemgraphKVStorage(MemgraphCrudMixin, BaseKVStorage):
                 if record["data"] is not None:
                     found[record["id"]] = record["data"]
             await result.consume()
-        return [self._decode(i, found.get(i)) for i in ids]
+        # LightRAG's own BaseKVStorage.get_by_ids declares -> list[dict[str, Any]], but its
+        # reference JsonKVStorage impl (and this one, to match) legitimately puts None in the
+        # slot of any id that doesn't exist -- the upstream annotation itself is imprecise.
+        return cast("list[dict[str, Any]]", [self._decode(i, found.get(i)) for i in ids])
 
     async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
         if not data:
@@ -126,13 +132,15 @@ class MemgraphKVStorage(MemgraphCrudMixin, BaseKVStorage):
         async with driver.session(database=get_database()) as session:
             await (
                 await session.run(
-                    f"""
-                    UNWIND $entries AS e
-                    MERGE (n:`{self._label}` {{id: e.id}})
-                    ON CREATE SET n.data = e.new_data, n.created_at = e.ts
-                    ON MATCH SET n.data = e.update_data
-                    SET n.updated_at = e.ts
-                    """,
+                    as_literal_query(
+                        f"""
+                        UNWIND $entries AS e
+                        MERGE (n:`{self._label}` {{id: e.id}})
+                        ON CREATE SET n.data = e.new_data, n.created_at = e.ts
+                        ON MATCH SET n.data = e.update_data
+                        SET n.updated_at = e.ts
+                        """
+                    ),
                     entries=entries,
                 )
             ).consume()

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/src/lightrag_memgraph/vector_impl.py

@@ -19,6 +19,7 @@ from lightrag.base import BaseVectorStorage
 from lightrag.utils import compute_mdhash_id, logger
 
 from ._connection import (
+    as_literal_query,
     close_driver,
     create_index,
     get_database,
@@ -130,18 +131,20 @@ class MemgraphVectorStorage(BaseVectorStorage):
             await self._ensure_vector_index(session)
             await (
                 await session.run(
-                    f"""
-                    UNWIND $entries AS e
-                    MERGE (n:`{self._label}` {{id: e.id}})
-                    ON CREATE SET n.created_at = e.ts
-                    SET n += e.props, n.embedding = e.embedding, n.updated_at = e.ts
-                    """,
+                    as_literal_query(
+                        f"""
+                        UNWIND $entries AS e
+                        MERGE (n:`{self._label}` {{id: e.id}})
+                        ON CREATE SET n.created_at = e.ts
+                        SET n += e.props, n.embedding = e.embedding, n.updated_at = e.ts
+                        """
+                    ),
                     entries=entries,
                 )
             ).consume()
         logger.debug(f"[{self.workspace}] Upserted {len(entries)} vectors to {self.namespace}")
 
-    async def query(self, query: str, top_k: int, query_embedding: list[float] = None) -> list[dict[str, Any]]:
+    async def query(self, query: str, top_k: int, query_embedding: list[float] | None = None) -> list[dict[str, Any]]:
         """Vector-search this namespace and return up to ``top_k`` live results.
 
         By design, Memgraph's native vector index operates at ``READ_UNCOMMITTED``
@@ -176,15 +179,17 @@ class MemgraphVectorStorage(BaseVectorStorage):
                 # dropping it.
                 # Index name is a sanitized string literal (not accepted as a query param).
                 result = await session.run(
-                    f"""
-                    CALL vector_search.search("{self._index_name}", $top_k, $embedding)
-                    YIELD node, similarity
-                    WITH collect({{node: node, similarity: similarity}}) AS cand_pairs
-                    UNWIND cand_pairs AS cp
-                    OPTIONAL MATCH (live:`{self._label}`) WHERE live = cp.node
-                    RETURN live.id AS id, cp.similarity AS similarity,
-                           CASE WHEN live IS NULL THEN NULL ELSE properties(live) END AS props
-                    """,
+                    as_literal_query(
+                        f"""
+                        CALL vector_search.search("{self._index_name}", $top_k, $embedding)
+                        YIELD node, similarity
+                        WITH collect({{node: node, similarity: similarity}}) AS cand_pairs
+                        UNWIND cand_pairs AS cp
+                        OPTIONAL MATCH (live:`{self._label}`) WHERE live = cp.node
+                        RETURN live.id AS id, cp.similarity AS similarity,
+                               CASE WHEN live IS NULL THEN NULL ELSE properties(live) END AS props
+                        """
+                    ),
                     top_k=int(top_k),
                     embedding=embedding,
                 )
@@ -225,7 +230,7 @@ class MemgraphVectorStorage(BaseVectorStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"MATCH (n:`{self._label}` {{id: $id}}) RETURN properties(n) AS props",
+                as_literal_query(f"MATCH (n:`{self._label}` {{id: $id}}) RETURN properties(n) AS props"),
                 id=id,
             )
             record = await result.single()
@@ -243,11 +248,13 @@ class MemgraphVectorStorage(BaseVectorStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"""
-                MATCH (n:`{self._label}`)
-                WHERE n.id IN $ids
-                RETURN n.id AS id, properties(n) AS props
-                """,
+                as_literal_query(
+                    f"""
+                    MATCH (n:`{self._label}`)
+                    WHERE n.id IN $ids
+                    RETURN n.id AS id, properties(n) AS props
+                    """
+                ),
                 ids=list(ids),
             )
             records = [record async for record in result]
@@ -266,11 +273,13 @@ class MemgraphVectorStorage(BaseVectorStorage):
         driver = await get_driver()
         async with driver.session(database=get_database(), default_access_mode="READ") as session:
             result = await session.run(
-                f"""
-                MATCH (n:`{self._label}`)
-                WHERE n.id IN $ids AND n.embedding IS NOT NULL
-                RETURN n.id AS id, n.embedding AS embedding
-                """,
+                as_literal_query(
+                    f"""
+                    MATCH (n:`{self._label}`)
+                    WHERE n.id IN $ids AND n.embedding IS NOT NULL
+                    RETURN n.id AS id, n.embedding AS embedding
+                    """
+                ),
                 ids=list(ids),
             )
             vectors = {record["id"]: list(record["embedding"]) async for record in result}
@@ -291,12 +300,14 @@ class MemgraphVectorStorage(BaseVectorStorage):
         async with driver.session(database=get_database()) as session:
             await (
                 await session.run(
-                    f"""
-                    MATCH (n:`{self._label}`)
-                    {where_clause}
-                    SET n.embedding = NULL
-                    DETACH DELETE n
-                    """,
+                    as_literal_query(
+                        f"""
+                        MATCH (n:`{self._label}`)
+                        {where_clause}
+                        SET n.embedding = NULL
+                        DETACH DELETE n
+                        """
+                    ),
                     **params,
                 )
             ).consume()
@@ -323,7 +334,7 @@ class MemgraphVectorStorage(BaseVectorStorage):
             driver = await get_driver()
             async with driver.session(database=get_database()) as session:
                 try:
-                    await (await session.run(f"DROP VECTOR INDEX {self._index_name}")).consume()
+                    await (await session.run(as_literal_query(f"DROP VECTOR INDEX {self._index_name}"))).consume()
                 except Exception as e:
                     logger.debug(f"[{self.workspace}] Dropping vector index {self._index_name} skipped: {e}")
             self._index_ready = False

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/tests/test_core.py

@@ -9,12 +9,13 @@ runs an insert/query).
 from __future__ import annotations
 
 import os
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
 from lightrag.utils import logger as lightrag_logger
 
-from lightrag_memgraph.core import _apply_lightrag_defaults, _bridge_lightrag_env_names
+from lightrag_memgraph.core import MemgraphLightRAGWrapper, _apply_lightrag_defaults, _bridge_lightrag_env_names
 from lightrag_memgraph.embeddings import memgraph_sentence_embed
 
 _MEMGRAPH_ENV_NAMES = ("MEMGRAPH_URL", "MEMGRAPH_URI", "MEMGRAPH_USER", "MEMGRAPH_USERNAME")
@@ -136,3 +137,49 @@ def test_no_openai_key_required_with_non_openai_llm_and_embedding_default(monkey
     kwargs = {"llm_model_func": lambda *a, **kw: None}
     _apply_lightrag_defaults(kwargs)  # must not raise
     assert kwargs["embedding_func"] is memgraph_sentence_embed
+
+
+@pytest.mark.asyncio
+async def test_afinalize_resets_shared_data_on_success():
+    wrapper = MemgraphLightRAGWrapper()
+    wrapper.rag = AsyncMock()
+
+    with patch("lightrag_memgraph.core.finalize_share_data") as mock_finalize_share_data:
+        await wrapper.afinalize()
+
+    wrapper.rag.finalize_storages.assert_awaited_once()
+    mock_finalize_share_data.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_afinalize_resets_shared_data_even_if_finalize_storages_raises():
+    """A transient error in finalize_storages() must not leave the
+    process-global shared-storage state stuck as initialized -- that would
+    break every subsequent LightRAG instance in this process, not just this
+    one's cleanup."""
+    wrapper = MemgraphLightRAGWrapper()
+    wrapper.rag = AsyncMock()
+    wrapper.rag.finalize_storages.side_effect = RuntimeError("transient network error")
+
+    with (
+        patch("lightrag_memgraph.core.finalize_share_data") as mock_finalize_share_data,
+        pytest.raises(RuntimeError, match="transient network error"),
+    ):
+        await wrapper.afinalize()
+
+    mock_finalize_share_data.assert_called_once()
+
+
+def test_workspace_reads_chunk_entity_relation_graph_workspace():
+    wrapper = MemgraphLightRAGWrapper()
+    wrapper.rag = MagicMock()
+    wrapper.rag.chunk_entity_relation_graph.workspace = "tenant-42"
+
+    assert wrapper.workspace == "tenant-42"
+
+
+def test_workspace_raises_before_initialize():
+    wrapper = MemgraphLightRAGWrapper()
+
+    with pytest.raises(RuntimeError, match="not initialized"):
+        _ = wrapper.workspace

{lightrag_memgraph-0.3.0 → lightrag_memgraph-0.3.2}/tests/test_integration_memgraph.py

@@ -35,6 +35,7 @@ from lightrag.utils import EmbeddingFunc
 
 from lightrag_memgraph import MemgraphLightRAGWrapper
 from lightrag_memgraph._connection import (
+    as_literal_query,
     close_driver,
     get_database,
     get_driver,
@@ -493,7 +494,8 @@ async def test_vector_query_logs_and_excludes_stale_candidates(
         # DETACH DELETE leaves a stale entry in the native vector index.
         driver = await get_driver()
         async with driver.session(database=get_database()) as session:
-            await (await session.run(f"MATCH (n:`{store._label}` {{id: 'v-cat'}}) DETACH DELETE n")).consume()
+            query = as_literal_query(f"MATCH (n:`{store._label}` {{id: 'v-cat'}}) DETACH DELETE n")
+            await (await session.run(query)).consume()
 
         # lightrag's logger has propagate=False, so caplog.at_level (which only
         # relies on root-logger propagation) never sees its records; attach the
@@ -612,9 +614,11 @@ async def test_wrapper_defaults_embedding_func_to_working_memgraph_sentence_embe
         # underlying callable, unwrapping the EmbeddingFunc nesting), so identity
         # doesn't survive the round trip -- the declared model/dim and, most
         # importantly, actual behaviour do.
-        assert rag.embedding_func.model_name == memgraph_sentence_embed.model_name
-        assert rag.embedding_func.embedding_dim == DEFAULT_EMBEDDING_DIM
-        vectors = await rag.embedding_func(["hello world", "graph databases are fast"])
+        embedding_func = rag.embedding_func
+        assert embedding_func is not None
+        assert embedding_func.model_name == memgraph_sentence_embed.model_name
+        assert embedding_func.embedding_dim == DEFAULT_EMBEDDING_DIM
+        vectors = await embedding_func(["hello world", "graph databases are fast"])
         assert vectors.shape == (2, DEFAULT_EMBEDDING_DIM)
     finally:
         if wrapper.rag is not None: