lightrag-memgraph 0.1.0__tar.gz

lightrag_memgraph-0.1.0/.gitignore

@@ -0,0 +1,2 @@
+lightrag_storage.out/
+lightrag.log

lightrag_memgraph-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Memgraph
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lightrag_memgraph-0.1.0/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.4
+Name: lightrag-memgraph
+Version: 0.1.0
+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
+Project-URL: Issues, https://github.com/memgraph/ai-toolkit/issues
+Author-email: Memgraph Tech Team <tech@memgraph.com>
+License: MIT
+License-File: LICENSE
+Keywords: graph,integration,lightrag,memgraph,toolkit
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.10
+Requires-Dist: lightrag-hku[api]==1.4.8.2
+Requires-Dist: memgraph-toolbox>=0.1.6
+Requires-Dist: numpy>=1.21.0
+Provides-Extra: dev
+Requires-Dist: black>=22.0.0; 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>=7.0.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# Why?
+
+* High-quality, slow and high-cost entity extraction.
+* Connect extracted entities with existing domain-specific entities.
+* Cost for your scale (e.g., $ per your document page).

lightrag_memgraph-0.1.0/README.md

@@ -0,0 +1,5 @@
+# Why?
+
+* High-quality, slow and high-cost entity extraction.
+* Connect extracted entities with existing domain-specific entities.
+* Cost for your scale (e.g., $ per your document page).

lightrag_memgraph-0.1.0/example.py

@@ -0,0 +1,101 @@
+import os
+import time
+import traceback
+
+import asyncio
+from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
+import shutil
+
+from lightrag_memgraph import MemgraphLightRAGWrapper
+from memgraph_toolbox.api.memgraph import Memgraph
+
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+WORKING_DIR = "./lightrag_storage.out"
+if os.path.exists(WORKING_DIR):
+    shutil.rmtree(WORKING_DIR)
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+memgraph = Memgraph()
+memgraph.query("MATCH (n) DETACH DELETE n;")
+
+DUMMY_TEXTS = [
+    """In the heart of the bustling city, a small bookstore stood as a
+    sanctuary for dreamers and thinkers alike. Its shelves, lined with
+    stories from every corner of the world, beckoned to those seeking
+    adventure, solace, or simply a quiet moment away from the noise.
+    The owner, an elderly gentleman with a gentle smile, greeted each
+    visitor as if they were an old friend. On rainy afternoons, the
+    soft patter of drops against the windows created a symphony that
+    mingled with the rustle of pages. Children gathered in the reading
+    nook, their imaginations ignited by tales of dragons and distant
+    lands. College students found refuge among the stacks, their minds
+    wandering as they prepared for exams. The bookstore was more than a
+    place to buy books; it was a haven where stories came alive,
+    friendships blossomed, and the magic of words wove its spell on all
+    who entered.""",
+    """Beneath the golden canopy of autumn leaves, a
+    quiet park unfolded its charm to those who wandered its winding
+    paths. Joggers traced familiar routes, their breath visible in the
+    crisp morning air, while elderly couples strolled hand in hand,
+    reminiscing about days gone by. Children’s laughter echoed from the
+    playground, where swings soared and slides became mountains to
+    conquer. A painter sat on a weathered bench, capturing the fiery
+    hues of the season on her canvas, her brush dancing with
+    inspiration. Nearby, a group of friends gathered for a picnic,
+    sharing stories and homemade treats as squirrels darted hopefully
+    around their feet. The gentle breeze carried the scent of earth and
+    fallen leaves, inviting all to pause and savor the moment. In this
+    tranquil oasis, time seemed to slow, offering a gentle reminder of
+    nature’s beauty and the simple joys that color everyday life.""",
+    """On the edge of a sleepy coastal village, a lighthouse stood
+    sentinel against the relentless waves. Its beacon, steadfast and
+    bright, guided fishermen safely home through fog and storm. The
+    keeper, a solitary figure with weathered hands, tended the light
+    with unwavering dedication, his days marked by the rhythm of tides
+    and the cries of gulls. Each evening, as the sun dipped below the
+    horizon, the village gathered on the shore to watch the sky ignite
+    in shades of orange and violet. Children chased the surf, their
+    laughter mingling with the roar of the sea. Local artisans
+    displayed their crafts at the market, their wares shaped by the
+    stories and traditions of generations. The lighthouse, a symbol of
+    hope and resilience, reminded all who saw it that even in the
+    darkest nights, a guiding light could be found, illuminating the
+    path home.""",
+]
+
+
+async def main():
+    lightrag_wrapper = MemgraphLightRAGWrapper(disable_embeddings=True)
+    try:
+        await lightrag_wrapper.initialize(
+            working_dir=WORKING_DIR,
+            max_parallel_insert=8,
+        )
+
+        total_time = 0.0
+        start_time = time.perf_counter()
+        await lightrag_wrapper.ainsert(
+            input=[text for text in DUMMY_TEXTS],
+            file_paths=[str(idx) for idx in range(len(DUMMY_TEXTS))],
+        )
+        end_time = time.perf_counter()
+        total_time += end_time - start_time
+        if len(DUMMY_TEXTS) > 0:
+            print(f"Average time per text: {total_time/len(DUMMY_TEXTS):.4f} seconds.")
+
+        rag = lightrag_wrapper.get_lightrag()
+        print(await rag.get_graph_labels())
+        kg_data = await rag.get_knowledge_graph(node_label="City", max_depth=3)
+        print("KNOWLEDGE GRAPH DATA:")
+        print(kg_data)
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        print(traceback.format_exc())
+    finally:
+        await lightrag_wrapper.afinalize()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

lightrag_memgraph-0.1.0/pyproject.toml

@@ -0,0 +1,68 @@
+[project]
+name = "lightrag-memgraph"
+version = "0.1.0"
+description = "LightRAG integration with Memgraph"
+readme = "README.md"
+requires-python = ">=3.10"
+
+authors = [
+    {name = "Memgraph Tech Team", email = "tech@memgraph.com"},
+]
+license = { text = "MIT" }
+keywords = ["memgraph", "graph", "integration", "toolkit", "lightrag"]
+
+classifiers = [
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+
+dependencies = [
+    "lightrag-hku[api]==1.4.8.2",
+    "memgraph-toolbox>=0.1.6",
+    "numpy>=1.21.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-asyncio>=0.21.0",
+    "black>=22.0.0",
+    "isort>=5.0.0",
+    "mypy>=1.0.0",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/memgraph/ai-toolkit"
+"Source" = "https://github.com/memgraph/ai-toolkit/tree/main/integrations/lightrag-memgraph"
+"Issues" = "https://github.com/memgraph/ai-toolkit/issues"
+
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+asyncio_mode = "strict"
+asyncio_default_fixture_loop_scope = "function"
+
+[build-system]
+requires = ["hatchling", "build", "setuptools"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/lightrag_memgraph"]
+
+[tool.black]
+line-length = 88
+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.1.0/src/lightrag_memgraph/__init__.py

@@ -0,0 +1,9 @@
+"""
+LightRAG integration with Memgraph.
+
+This package provides a wrapper around LightRAG that uses Memgraph as the graph storage backend.
+"""
+
+from .core import MemgraphLightRAGWrapper
+
+__all__ = ["MemgraphLightRAGWrapper"]

lightrag_memgraph-0.1.0/src/lightrag_memgraph/core.py

@@ -0,0 +1,85 @@
+import os
+from typing import Optional
+import logging
+
+from lightrag import LightRAG
+from lightrag.kg.shared_storage import initialize_pipeline_status
+from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
+from lightrag.utils import setup_logger
+import numpy as np
+
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+MEMGRAPH_URI = os.getenv("MEMGRAPH_URI", "bolt://localhost:7687")
+os.environ["MEMGRAPH_URI"] = MEMGRAPH_URI
+
+
+class DummyEmbed:
+    def __init__(self, dim: int = 1):
+        self.embedding_dim = dim
+
+    async def __call__(self, texts: list[str]) -> np.ndarray:
+        return np.ones((len(texts), self.embedding_dim), dtype=float)
+
+
+class MemgraphLightRAGWrapper:
+    def __init__(
+        self,
+        log_level: str = "INFO",
+        disable_embeddings: bool = False,
+    ):
+        self.log_level = log_level
+        self.disable_embeddings = disable_embeddings
+        self.rag: Optional[LightRAG] = None
+
+    # https://github.com/HKUDS/LightRAG/blob/main/lightrag/lightrag.py
+    # https://github.com/HKUDS/LightRAG/blob/main/lightrag/llm
+    async def initialize(self, **lightrag_kwargs) -> None:
+        setup_logger("lightrag", level=self.log_level)
+        logging.getLogger("nano-vectordb").setLevel(self.log_level)
+        logging.getLogger("pikepdf").setLevel(self.log_level)
+        if self.disable_embeddings:
+            lightrag_kwargs["embedding_func"] = DummyEmbed(dim=1)
+            lightrag_kwargs["vector_storage"] = "NanoVectorDBStorage"
+        if "working_dir" in lightrag_kwargs:
+            working_dir = lightrag_kwargs["working_dir"]
+            if not os.path.exists(working_dir):
+                os.mkdir(working_dir)
+        if "llm_model_func" not in lightrag_kwargs:
+            lightrag_kwargs["llm_model_func"] = gpt_4o_mini_complete
+        if "embedding_func" not in lightrag_kwargs:
+            lightrag_kwargs["embedding_func"] = openai_embed
+        if (
+            lightrag_kwargs["llm_model_func"] == gpt_4o_mini_complete
+            or lightrag_kwargs["embedding_func"] == openai_embed
+        ):
+            openai_api_key = os.environ.get("OPENAI_API_KEY")
+            if not openai_api_key:
+                raise EnvironmentError(
+                    "OPENAI_API_KEY environment variable is not set. Please set your OpenAI API key."
+                )
+        self.rag = LightRAG(graph_storage="MemgraphStorage", **lightrag_kwargs)
+        await self.rag.initialize_storages()
+        await initialize_pipeline_status()
+
+    def get_lightrag(self) -> LightRAG:
+        if self.rag is None:
+            raise RuntimeError("LightRAG not initialized. Call initialize() first.")
+        return self.rag
+
+    # https://github.com/HKUDS/LightRAG/blob/main/lightrag/lightrag.py
+    async def ainsert(self, **kwargs) -> None:
+        """
+        Example call: await lightrag_wrapper.ainsert(input=text, file_paths=[id])
+
+        If you want to inject info under each entity about the source input,
+        pass file_paths as a list of strings (ids don't work, not written under each entity).
+        """
+        if self.rag is None:
+            raise RuntimeError("LightRAG not initialized. Call initialize() first.")
+        await self.rag.ainsert(**kwargs)
+
+    async def afinalize(self) -> None:
+        if self.rag is None:
+            raise RuntimeError("LightRAG not initialized. Call initialize() first.")
+        await self.rag.finalize_storages()