lightrag-hku 0.0.1__tar.gz

lightrag_hku-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Gustavo Ye
+
+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_hku-0.0.1/PKG-INFO

@@ -0,0 +1,39 @@
+Metadata-Version: 2.1
+Name: lightrag-hku
+Version: 0.0.1
+Summary: LightRAG: Simple and Fast Retrieval-Augmented Generation
+Home-page: https://github.com/HKUDS/GraphEdit
+Author: ZiruiGuo
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: openai
+Requires-Dist: tiktoken
+Requires-Dist: networkx
+Requires-Dist: graspologic
+Requires-Dist: nano-vectordb
+Requires-Dist: hnswlib
+Requires-Dist: xxhash
+Requires-Dist: tenacity
+
+# LightRAG: Simple and Fast Retrieval-Augmented Generation
+
+## Citation
+
+```
+@article{guo2024lightrag,
+title={LightRAG: Simple and Fast Retrieval-Augmented Generation}, 
+author={Zirui Guo and Lianghao Xia and Yanhua Yu and Tu Ao and Chao Huang},
+year={2024},
+eprint={},
+archivePrefix={arXiv},
+primaryClass={cs.IR}
+}
+```
+
+## Acknowledgement
+
+The structure of this code is based on [nano-graphrag](https://github.com/gusye1234/nano-graphrag).

lightrag_hku-0.0.1/README.md

@@ -0,0 +1,18 @@
+# LightRAG: Simple and Fast Retrieval-Augmented Generation
+
+## Citation
+
+```
+@article{guo2024lightrag,
+title={LightRAG: Simple and Fast Retrieval-Augmented Generation}, 
+author={Zirui Guo and Lianghao Xia and Yanhua Yu and Tu Ao and Chao Huang},
+year={2024},
+eprint={},
+archivePrefix={arXiv},
+primaryClass={cs.IR}
+}
+```
+
+## Acknowledgement
+
+The structure of this code is based on [nano-graphrag](https://github.com/gusye1234/nano-graphrag).

lightrag_hku-0.0.1/lightrag/__init__.py

@@ -0,0 +1,5 @@
+from .lightrag import LightRAG, QueryParam
+
+__version__ = "0.0.1"
+__author__ = "Zirui Guo"
+__url__ = "https://github.com/HKUDS/GraphEdit"

lightrag_hku-0.0.1/lightrag/base.py

@@ -0,0 +1,116 @@
+from dataclasses import dataclass, field
+from typing import TypedDict, Union, Literal, Generic, TypeVar
+
+import numpy as np
+
+from .utils import EmbeddingFunc
+
+TextChunkSchema = TypedDict(
+    "TextChunkSchema",
+    {"tokens": int, "content": str, "full_doc_id": str, "chunk_order_index": int},
+)
+
+T = TypeVar("T")
+
+@dataclass
+class QueryParam:
+    mode: Literal["local", "global", "hybird", "naive"] = "global"
+    only_need_context: bool = False
+    response_type: str = "Multiple Paragraphs"
+    top_k: int = 60
+    max_token_for_text_unit: int = 4000 
+    max_token_for_global_context: int = 4000
+    max_token_for_local_context: int = 4000 
+
+
+@dataclass
+class StorageNameSpace:
+    namespace: str
+    global_config: dict
+
+    async def index_done_callback(self):
+        """commit the storage operations after indexing"""
+        pass
+
+    async def query_done_callback(self):
+        """commit the storage operations after querying"""
+        pass
+
+@dataclass
+class BaseVectorStorage(StorageNameSpace):
+    embedding_func: EmbeddingFunc
+    meta_fields: set = field(default_factory=set)
+
+    async def query(self, query: str, top_k: int) -> list[dict]:
+        raise NotImplementedError
+
+    async def upsert(self, data: dict[str, dict]):
+        """Use 'content' field from value for embedding, use key as id.
+        If embedding_func is None, use 'embedding' field from value
+        """
+        raise NotImplementedError
+
+@dataclass
+class BaseKVStorage(Generic[T], StorageNameSpace):
+    async def all_keys(self) -> list[str]:
+        raise NotImplementedError
+
+    async def get_by_id(self, id: str) -> Union[T, None]:
+        raise NotImplementedError
+
+    async def get_by_ids(
+        self, ids: list[str], fields: Union[set[str], None] = None
+    ) -> list[Union[T, None]]:
+        raise NotImplementedError
+
+    async def filter_keys(self, data: list[str]) -> set[str]:
+        """return un-exist keys"""
+        raise NotImplementedError
+
+    async def upsert(self, data: dict[str, T]):
+        raise NotImplementedError
+
+    async def drop(self):
+        raise NotImplementedError
+    
+
+@dataclass
+class BaseGraphStorage(StorageNameSpace):
+    async def has_node(self, node_id: str) -> bool:
+        raise NotImplementedError
+
+    async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
+        raise NotImplementedError
+
+    async def node_degree(self, node_id: str) -> int:
+        raise NotImplementedError
+
+    async def edge_degree(self, src_id: str, tgt_id: str) -> int:
+        raise NotImplementedError
+
+    async def get_node(self, node_id: str) -> Union[dict, None]:
+        raise NotImplementedError
+
+    async def get_edge(
+        self, source_node_id: str, target_node_id: str
+    ) -> Union[dict, None]:
+        raise NotImplementedError
+
+    async def get_node_edges(
+        self, source_node_id: str
+    ) -> Union[list[tuple[str, str]], None]:
+        raise NotImplementedError
+
+    async def upsert_node(self, node_id: str, node_data: dict[str, str]):
+        raise NotImplementedError
+
+    async def upsert_edge(
+        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
+    ):
+        raise NotImplementedError
+
+    async def clustering(self, algorithm: str):
+        raise NotImplementedError
+
+    async def embed_nodes(self, algorithm: str) -> tuple[np.ndarray, list[str]]:
+        raise NotImplementedError("Node embedding is not used in lightrag.")

lightrag_hku-0.0.1/lightrag/lightrag.py

@@ -0,0 +1,289 @@
+import asyncio
+import os
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from functools import partial
+from typing import Type, cast
+
+from .llm import gpt_4o_complete, gpt_4o_mini_complete, openai_embedding
+from .operate import (
+    chunking_by_token_size,
+    extract_entities,
+    local_query,
+    global_query,
+    hybird_query,
+    naive_query,
+)
+
+from .storage import (
+    JsonKVStorage,
+    NanoVectorDBStorage,
+    NetworkXStorage,
+)
+from .utils import (
+    EmbeddingFunc,
+    compute_mdhash_id,
+    limit_async_func_call,
+    convert_response_to_json,
+    logger,
+    set_logger,
+)
+from .base import (
+    BaseGraphStorage,
+    BaseKVStorage,
+    BaseVectorStorage,
+    StorageNameSpace,
+    QueryParam,
+)
+
+def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
+    try:
+        # If there is already an event loop, use it.
+        loop = asyncio.get_event_loop()
+    except RuntimeError:
+        # If in a sub-thread, create a new event loop.
+        logger.info("Creating a new event loop in a sub-thread.")
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+    return loop
+
+@dataclass
+class LightRAG:
+    working_dir: str = field(
+        default_factory=lambda: f"./lightrag_cache_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}"
+    )
+
+    # text chunking
+    chunk_token_size: int = 1200
+    chunk_overlap_token_size: int = 100
+    tiktoken_model_name: str = "gpt-4o-mini"
+
+    # entity extraction
+    entity_extract_max_gleaning: int = 1
+    entity_summary_to_max_tokens: int = 500
+
+    # node embedding
+    node_embedding_algorithm: str = "node2vec"
+    node2vec_params: dict = field(
+        default_factory=lambda: {
+            "dimensions": 1536,
+            "num_walks": 10,
+            "walk_length": 40,
+            "num_walks": 10,
+            "window_size": 2,
+            "iterations": 3,
+            "random_seed": 3,
+        }
+    )
+
+    # text embedding
+    embedding_func: EmbeddingFunc = field(default_factory=lambda: openai_embedding)
+    embedding_batch_num: int = 32
+    embedding_func_max_async: int = 16
+
+    # LLM
+    llm_model_func: callable = gpt_4o_mini_complete
+    llm_model_max_token_size: int = 32768
+    llm_model_max_async: int = 16
+
+    # storage
+    key_string_value_json_storage_cls: Type[BaseKVStorage] = JsonKVStorage
+    vector_db_storage_cls: Type[BaseVectorStorage] = NanoVectorDBStorage
+    vector_db_storage_cls_kwargs: dict = field(default_factory=dict)
+    graph_storage_cls: Type[BaseGraphStorage] = NetworkXStorage
+    enable_llm_cache: bool = True
+
+    # extension
+    addon_params: dict = field(default_factory=dict)
+    convert_response_to_json_func: callable = convert_response_to_json
+
+    def __post_init__(self):        
+        log_file = os.path.join(self.working_dir, "lightrag.log")
+        set_logger(log_file)
+        logger.info(f"Logger initialized for working directory: {self.working_dir}")
+        
+        _print_config = ",\n  ".join([f"{k} = {v}" for k, v in asdict(self).items()])
+        logger.debug(f"LightRAG init with param:\n  {_print_config}\n")
+
+        if not os.path.exists(self.working_dir):
+            logger.info(f"Creating working directory {self.working_dir}")
+            os.makedirs(self.working_dir)
+
+        self.full_docs = self.key_string_value_json_storage_cls(
+            namespace="full_docs", global_config=asdict(self)
+        )
+
+        self.text_chunks = self.key_string_value_json_storage_cls(
+            namespace="text_chunks", global_config=asdict(self)
+        )
+
+        self.llm_response_cache = (
+            self.key_string_value_json_storage_cls(
+                namespace="llm_response_cache", global_config=asdict(self)
+            )
+            if self.enable_llm_cache
+            else None
+        )
+        self.chunk_entity_relation_graph = self.graph_storage_cls(
+            namespace="chunk_entity_relation", global_config=asdict(self)
+        )
+        self.embedding_func = limit_async_func_call(self.embedding_func_max_async)(
+            self.embedding_func
+        )
+        self.entities_vdb = (
+            self.vector_db_storage_cls(
+                namespace="entities",
+                global_config=asdict(self),
+                embedding_func=self.embedding_func,
+                meta_fields={"entity_name"}
+            )
+        )
+        self.relationships_vdb = (
+            self.vector_db_storage_cls(
+                namespace="relationships",
+                global_config=asdict(self),
+                embedding_func=self.embedding_func,
+                meta_fields={"src_id", "tgt_id"}
+            )
+        )
+        self.chunks_vdb = (
+            self.vector_db_storage_cls(
+                namespace="chunks",
+                global_config=asdict(self),
+                embedding_func=self.embedding_func,
+            )
+        )
+        
+        self.llm_model_func = limit_async_func_call(self.llm_model_max_async)(
+            partial(self.llm_model_func, hashing_kv=self.llm_response_cache)
+        )
+
+    def insert(self, string_or_strings):
+        loop = always_get_an_event_loop()
+        return loop.run_until_complete(self.ainsert(string_or_strings))
+
+    async def ainsert(self, string_or_strings):
+        try:
+            if isinstance(string_or_strings, str):
+                string_or_strings = [string_or_strings]
+            # ---------- new docs
+            new_docs = {
+                compute_mdhash_id(c.strip(), prefix="doc-"): {"content": c.strip()}
+                for c in string_or_strings
+            }
+            _add_doc_keys = await self.full_docs.filter_keys(list(new_docs.keys()))
+            new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}
+            if not len(new_docs):
+                logger.warning(f"All docs are already in the storage")
+                return
+            logger.info(f"[New Docs] inserting {len(new_docs)} docs")
+
+            # ---------- chunking
+            inserting_chunks = {}
+            for doc_key, doc in new_docs.items():
+                chunks = {
+                    compute_mdhash_id(dp["content"], prefix="chunk-"): {
+                        **dp,
+                        "full_doc_id": doc_key,
+                    }
+                    for dp in chunking_by_token_size(
+                        doc["content"],
+                        overlap_token_size=self.chunk_overlap_token_size,
+                        max_token_size=self.chunk_token_size,
+                        tiktoken_model=self.tiktoken_model_name,
+                    )
+                }
+                inserting_chunks.update(chunks)
+            _add_chunk_keys = await self.text_chunks.filter_keys(
+                list(inserting_chunks.keys())
+            )
+            inserting_chunks = {
+                k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
+            }
+            if not len(inserting_chunks):
+                logger.warning(f"All chunks are already in the storage")
+                return
+            logger.info(f"[New Chunks] inserting {len(inserting_chunks)} chunks")
+
+            await self.chunks_vdb.upsert(inserting_chunks)
+
+            # ---------- commit upsertings and indexing
+            await self.full_docs.upsert(new_docs)
+            await self.text_chunks.upsert(inserting_chunks)
+        finally:
+            await self._insert_done()
+
+    async def _insert_done(self):
+        tasks = []
+        for storage_inst in [
+            self.full_docs,
+            self.text_chunks,
+            self.llm_response_cache,
+            self.entities_vdb,
+            self.relationships_vdb,
+            self.chunks_vdb,
+            self.chunk_entity_relation_graph,
+        ]:
+            if storage_inst is None:
+                continue
+            tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
+        await asyncio.gather(*tasks)
+
+    def query(self, query: str, param: QueryParam = QueryParam()):
+        loop = always_get_an_event_loop()
+        return loop.run_until_complete(self.aquery(query, param))
+    
+    async def aquery(self, query: str, param: QueryParam = QueryParam()):
+        if param.mode == "local":
+            response = await local_query(
+                query,
+                self.chunk_entity_relation_graph,
+                self.entities_vdb,
+                self.relationships_vdb,
+                self.text_chunks,
+                param,
+                asdict(self),
+            )
+        elif param.mode == "global":
+            response = await global_query(
+                query,
+                self.chunk_entity_relation_graph,
+                self.entities_vdb,
+                self.relationships_vdb,
+                self.text_chunks,
+                param,
+                asdict(self),
+            )
+        elif param.mode == "hybird":
+            response = await hybird_query(
+                query,
+                self.chunk_entity_relation_graph,
+                self.entities_vdb,
+                self.relationships_vdb,
+                self.text_chunks,
+                param,
+                asdict(self),
+            )
+        elif param.mode == "naive":
+            response = await naive_query(
+                query,
+                self.chunks_vdb,
+                self.text_chunks,
+                param,
+                asdict(self),
+            )
+        else:
+            raise ValueError(f"Unknown mode {param.mode}")
+        await self._query_done()
+        return response
+    
+
+    async def _query_done(self):
+        tasks = []
+        for storage_inst in [self.llm_response_cache]:
+            if storage_inst is None:
+                continue
+            tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
+        await asyncio.gather(*tasks)
+
+    

lightrag_hku-0.0.1/lightrag/llm.py

@@ -0,0 +1,88 @@
+import os
+import numpy as np
+from openai import AsyncOpenAI, APIConnectionError, RateLimitError, Timeout
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    retry_if_exception_type,
+)
+
+from .base import BaseKVStorage
+from .utils import compute_args_hash, wrap_embedding_func_with_attrs
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def openai_complete_if_cache(
+    model, prompt, system_prompt=None, history_messages=[], **kwargs
+) -> str:
+    openai_async_client = AsyncOpenAI()
+    hashing_kv: BaseKVStorage = kwargs.pop("hashing_kv", None)
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+    if hashing_kv is not None:
+        args_hash = compute_args_hash(model, messages)
+        if_cache_return = await hashing_kv.get_by_id(args_hash)
+        if if_cache_return is not None:
+            return if_cache_return["return"]
+
+    response = await openai_async_client.chat.completions.create(
+        model=model, messages=messages, **kwargs
+    )
+
+    if hashing_kv is not None:
+        await hashing_kv.upsert(
+            {args_hash: {"return": response.choices[0].message.content, "model": model}}
+        )
+    return response.choices[0].message.content
+
+async def gpt_4o_complete(
+    prompt, system_prompt=None, history_messages=[], **kwargs
+) -> str:
+    return await openai_complete_if_cache(
+        "gpt-4o",
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        **kwargs,
+    )
+
+
+async def gpt_4o_mini_complete(
+    prompt, system_prompt=None, history_messages=[], **kwargs
+) -> str:
+    return await openai_complete_if_cache(
+        "gpt-4o-mini",
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        **kwargs,
+    )
+
+@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def openai_embedding(texts: list[str]) -> np.ndarray:
+    openai_async_client = AsyncOpenAI()
+    response = await openai_async_client.embeddings.create(
+        model="text-embedding-3-small", input=texts, encoding_format="float"
+    )
+    return np.array([dp.embedding for dp in response.data])
+
+if __name__ == "__main__":
+    import asyncio
+
+    async def main():
+        result = await gpt_4o_mini_complete('How are you?')
+        print(result)
+
+    asyncio.run(main())