Merge pull request #1 from HKUDS/main

pull head

README.md

@@ -594,7 +594,7 @@ if __name__ == "__main__":
 | **llm\_model\_kwargs** | `dict` | Additional parameters for LLM generation |     |
 | **vector\_db\_storage\_cls\_kwargs** | `dict` | Additional parameters for vector database (currently not used) |     |
 | **enable\_llm\_cache** | `bool` | If `TRUE`, stores LLM results in cache; repeated prompts return cached responses | `TRUE` |
-| **addon\_params** | `dict` | Additional parameters, e.g., `{"example_number": 1, "language": "Simplified Chinese"}`: sets example limit and output language | `example_number: all examples, language: English` |
+| **addon\_params** | `dict` | Additional parameters, e.g., `{"example_number": 1, "language": "Simplified Chinese", "entity_types": ["organization", "person", "geo", "event"]}`: sets example limit and output language | `example_number: all examples, language: English` |
 | **convert\_response\_to\_json\_func** | `callable` | Not used | `convert_response_to_json` |
 | **embedding\_cache\_config** | `dict` | Configuration for question-answer caching. Contains three parameters:<br>- `enabled`: Boolean value to enable/disable cache lookup functionality. When enabled, the system will check cached responses before generating new answers.<br>- `similarity_threshold`: Float value (0-1), similarity threshold. When a new question's similarity with a cached question exceeds this threshold, the cached answer will be returned directly without calling the LLM.<br>- `use_llm_check`: Boolean value to enable/disable LLM similarity verification. When enabled, LLM will be used as a secondary check to verify the similarity between questions before returning cached answers. | Default: `{"enabled": False, "similarity_threshold": 0.95, "use_llm_check": False}` |
 

examples/lightrag_ollama_age_demo.py

@@ -0,0 +1,80 @@
+import asyncio
+import inspect
+import logging
+import os
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm import ollama_embedding, ollama_model_complete
+from lightrag.utils import EmbeddingFunc
+
+WORKING_DIR = "./dickens_age"
+
+logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+# AGE
+os.environ["AGE_POSTGRES_DB"] = "postgresDB"
+os.environ["AGE_POSTGRES_USER"] = "postgresUser"
+os.environ["AGE_POSTGRES_PASSWORD"] = "postgresPW"
+os.environ["AGE_POSTGRES_HOST"] = "localhost"
+os.environ["AGE_POSTGRES_PORT"] = "5455"
+os.environ["AGE_GRAPH_NAME"] = "dickens"
+
+rag = LightRAG(
+    working_dir=WORKING_DIR,
+    llm_model_func=ollama_model_complete,
+    llm_model_name="llama3.1:8b",
+    llm_model_max_async=4,
+    llm_model_max_token_size=32768,
+    llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 32768}},
+    embedding_func=EmbeddingFunc(
+        embedding_dim=768,
+        max_token_size=8192,
+        func=lambda texts: ollama_embedding(
+            texts, embed_model="nomic-embed-text", host="http://localhost:11434"
+        ),
+    ),
+    graph_storage="AGEStorage",
+)
+
+with open("./book.txt", "r", encoding="utf-8") as f:
+    rag.insert(f.read())
+
+# Perform naive search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
+)
+
+# Perform local search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+)
+
+# Perform global search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+)
+
+# Perform hybrid search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
+)
+
+# stream response
+resp = rag.query(
+    "What are the top themes in this story?",
+    param=QueryParam(mode="hybrid", stream=True),
+)
+
+
+async def print_stream(stream):
+    async for chunk in stream:
+        print(chunk, end="", flush=True)
+
+
+if inspect.isasyncgen(resp):
+    asyncio.run(print_stream(resp))
+else:
+    print(resp)

examples/lightrag_tidb_demo.py

@@ -0,0 +1,127 @@
+import asyncio
+import os
+
+import numpy as np
+
+from lightrag import LightRAG, QueryParam
+from lightrag.kg.tidb_impl import TiDB
+from lightrag.llm import siliconcloud_embedding, openai_complete_if_cache
+from lightrag.utils import EmbeddingFunc
+
+WORKING_DIR = "./dickens"
+
+# We use SiliconCloud API to call LLM on Oracle Cloud
+# More docs here https://docs.siliconflow.cn/introduction
+BASE_URL = "https://api.siliconflow.cn/v1/"
+APIKEY = ""
+CHATMODEL = ""
+EMBEDMODEL = ""
+
+TIDB_HOST = ""
+TIDB_PORT = ""
+TIDB_USER = ""
+TIDB_PASSWORD = ""
+TIDB_DATABASE = ""
+
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+async def llm_model_func(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    return await openai_complete_if_cache(
+        CHATMODEL,
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=APIKEY,
+        base_url=BASE_URL,
+        **kwargs,
+    )
+
+
+async def embedding_func(texts: list[str]) -> np.ndarray:
+    return await siliconcloud_embedding(
+        texts,
+        # model=EMBEDMODEL,
+        api_key=APIKEY,
+    )
+
+
+async def get_embedding_dim():
+    test_text = ["This is a test sentence."]
+    embedding = await embedding_func(test_text)
+    embedding_dim = embedding.shape[1]
+    return embedding_dim
+
+
+async def main():
+    try:
+        # Detect embedding dimension
+        embedding_dimension = await get_embedding_dim()
+        print(f"Detected embedding dimension: {embedding_dimension}")
+
+        # Create TiDB DB connection
+        tidb = TiDB(
+            config={
+                "host": TIDB_HOST,
+                "port": TIDB_PORT,
+                "user": TIDB_USER,
+                "password": TIDB_PASSWORD,
+                "database": TIDB_DATABASE,
+                "workspace": "company",  # specify which docs you want to store and query
+            }
+        )
+
+        # Check if TiDB DB tables exist, if not, tables will be created
+        await tidb.check_tables()
+
+        # Initialize LightRAG
+        # We use TiDB DB as the KV/vector
+        # You can add `addon_params={"example_number": 1, "language": "Simplfied Chinese"}` to control the prompt
+        rag = LightRAG(
+            enable_llm_cache=False,
+            working_dir=WORKING_DIR,
+            chunk_token_size=512,
+            llm_model_func=llm_model_func,
+            embedding_func=EmbeddingFunc(
+                embedding_dim=embedding_dimension,
+                max_token_size=512,
+                func=embedding_func,
+            ),
+            kv_storage="TiDBKVStorage",
+            vector_storage="TiDBVectorDBStorage",
+        )
+
+        if rag.llm_response_cache:
+            rag.llm_response_cache.db = tidb
+        rag.full_docs.db = tidb
+        rag.text_chunks.db = tidb
+        rag.entities_vdb.db = tidb
+        rag.relationships_vdb.db = tidb
+        rag.chunks_vdb.db = tidb
+
+        # Extract and Insert into LightRAG storage
+        with open("./dickens/demo.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+
+        # Perform search in different modes
+        modes = ["naive", "local", "global", "hybrid"]
+        for mode in modes:
+            print("=" * 20, mode, "=" * 20)
+            print(
+                await rag.aquery(
+                    "What are the top themes in this story?",
+                    param=QueryParam(mode=mode),
+                )
+            )
+            print("-" * 100, "\n")
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/lightrag_zhipu_demo.py

@@ -0,0 +1,55 @@
+import os
+import logging
+
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm import zhipu_complete, zhipu_embedding
+from lightrag.utils import EmbeddingFunc
+
+WORKING_DIR = "./dickens"
+
+logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+api_key = os.environ.get("ZHIPUAI_API_KEY")
+if api_key is None:
+    raise Exception("Please set ZHIPU_API_KEY in your environment")
+
+
+rag = LightRAG(
+    working_dir=WORKING_DIR,
+    llm_model_func=zhipu_complete,
+    llm_model_name="glm-4-flashx",  # Using the most cost/performance balance model, but you can change it here.
+    llm_model_max_async=4,
+    llm_model_max_token_size=32768,
+    embedding_func=EmbeddingFunc(
+        embedding_dim=2048,  # Zhipu embedding-3 dimension
+        max_token_size=8192,
+        func=lambda texts: zhipu_embedding(texts),
+    ),
+)
+
+with open("./book.txt", "r", encoding="utf-8") as f:
+    rag.insert(f.read())
+
+# Perform naive search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
+)
+
+# Perform local search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+)
+
+# Perform global search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+)
+
+# Perform hybrid search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
+)

lightrag/__init__.py

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