lightrag-memgraph 0.1.0__py3-none-any.whl

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/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()

lightrag_memgraph-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,6 @@
+lightrag_memgraph/__init__.py,sha256=-bcTvh8I881iZnvB55HX8JQ3VAbyvVvlsvRn6i8x4no,224
+lightrag_memgraph/core.py,sha256=pu57AZAcFZptbUFhGMttrs7BlJMTKf0gT2J5vVjUdz4,3430
+lightrag_memgraph-0.1.0.dist-info/METADATA,sha256=sbF8_i7Ja8uKBR_U8kvgrjFIna_WC1X1BPgeReX0q90,1437
+lightrag_memgraph-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+lightrag_memgraph-0.1.0.dist-info/licenses/LICENSE,sha256=Zs4MZAXqusfO2KGQoAOuIqEUtojbyIYIkiJQ0BTt45o,1065
+lightrag_memgraph-0.1.0.dist-info/RECORD,,

lightrag_memgraph-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

lightrag_memgraph-0.1.0.dist-info/licenses/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.