Merge branch 'HKUDS:main' into main

README.md

@@ -12,7 +12,7 @@
     </p>
      <p>
           <img src='https://img.shields.io/github/stars/hkuds/lightrag?color=green&style=social' />
-        <img src="https://img.shields.io/badge/python->=3.10-blue">
+        <img src="https://img.shields.io/badge/python-3.10-blue">
         <a href="https://pypi.org/project/lightrag-hku/"><img src="https://img.shields.io/pypi/v/lightrag-hku.svg"></a>
         <a href="https://pepy.tech/project/lightrag-hku"><img src="https://static.pepy.tech/badge/lightrag-hku/month"></a>
     </p>
@@ -26,7 +26,8 @@ This repository hosts the code of LightRAG. The structure of this code is based
 </div>
 
 ## 🎉 News
-- [x] [2025.01.06]🎯📢LightRAG now supports [PostgreSQL for Storage](https://github.com/HKUDS/LightRAG?tab=readme-ov-file#using-postgres-for-storage).
+- [x] [2025.01.13]🎯📢Our team has released [MiniRAG](https://github.com/HKUDS/MiniRAG) making RAG simpler with small models.
+- [x] [2025.01.06]🎯📢You can now [use PostgreSQL for Storage](#using-postgresql-for-storage).
 - [x] [2024.12.31]🎯📢LightRAG now supports [deletion by document ID](https://github.com/HKUDS/LightRAG?tab=readme-ov-file#delete).
 - [x] [2024.11.25]🎯📢LightRAG now supports seamless integration of [custom knowledge graphs](https://github.com/HKUDS/LightRAG?tab=readme-ov-file#insert-custom-kg), empowering users to enhance the system with their own domain expertise.
 - [x] [2024.11.19]🎯📢A comprehensive guide to LightRAG is now available on [LearnOpenCV](https://learnopencv.com/lightrag). Many thanks to the blog author.
@@ -361,6 +362,18 @@ see test_neo4j.py for a working example.
 For production level scenarios you will most likely want to leverage an enterprise solution. PostgreSQL can provide a one-stop solution for you as KV store, VectorDB (pgvector) and GraphDB (apache AGE).
 * PostgreSQL is lightweight,the whole binary distribution including all necessary plugins can be zipped to 40MB: Ref to [Windows Release](https://github.com/ShanGor/apache-age-windows/releases/tag/PG17%2Fv1.5.0-rc0) as it is easy to install for Linux/Mac.
 * How to start? Ref to: [examples/lightrag_zhipu_postgres_demo.py](https://github.com/HKUDS/LightRAG/blob/main/examples/lightrag_zhipu_postgres_demo.py)
+* Create index for AGE example: (Change below `dickens` to your graph name if necessary)
+  ```
+  SET search_path = ag_catalog, "$user", public;
+  CREATE INDEX idx_entity ON dickens."Entity" USING gin (agtype_access_operator(properties, '"node_id"'));
+  ```
+* Known issue of the Apache AGE: The released versions got below issue:
+  > You might find that the properties of the nodes/edges are empty.
+  > It is a known issue of the release version: https://github.com/apache/age/pull/1721
+  >
+  > You can Compile the AGE from source code and fix it.
+
+
 
 ### Insert Custom KG
 
@@ -912,12 +925,14 @@ pip install -e ".[api]"
 
 ### Prerequisites
 
-Before running any of the servers, ensure you have the corresponding backend service running:
+Before running any of the servers, ensure you have the corresponding backend service running for both llm and embedding.
+The new api allows you to mix different bindings for llm/embeddings.
+For example, you have the possibility to use ollama for the embedding and openai for the llm.
 
 #### For LoLLMs Server
 - LoLLMs must be running and accessible
 - Default connection: http://localhost:9600
-- Configure using --lollms-host if running on a different host/port
+- Configure using --llm-binding-host and/or --embedding-binding-host if running on a different host/port
 
 #### For Ollama Server
 - Ollama must be running and accessible
@@ -953,113 +968,96 @@ The output of the last command will give you the endpoint and the key for the Op
 
 Each server has its own specific configuration options:
 
-#### LoLLMs Server Options
+#### LightRag Server Options
 
 | Parameter | Default | Description |
 |-----------|---------|-------------|
-| --host | 0.0.0.0 | RAG server host |
-| --port | 9621 | RAG server port |
-| --model | mistral-nemo:latest | LLM model name |
+| --host | 0.0.0.0 | Server host |
+| --port | 9621 | Server port |
+| --llm-binding | ollama | LLM binding to be used. Supported: lollms, ollama, openai |
+| --llm-binding-host | (dynamic) | LLM server host URL. Defaults based on binding: http://localhost:11434 (ollama), http://localhost:9600 (lollms), https://api.openai.com/v1 (openai) |
+| --llm-model | mistral-nemo:latest | LLM model name |
+| --embedding-binding | ollama | Embedding binding to be used. Supported: lollms, ollama, openai |
+| --embedding-binding-host | (dynamic) | Embedding server host URL. Defaults based on binding: http://localhost:11434 (ollama), http://localhost:9600 (lollms), https://api.openai.com/v1 (openai) |
 | --embedding-model | bge-m3:latest | Embedding model name |
-| --lollms-host | http://localhost:9600 | LoLLMS backend URL |
-| --working-dir | ./rag_storage | Working directory for RAG |
+| --working-dir | ./rag_storage | Working directory for RAG storage |
+| --input-dir | ./inputs | Directory containing input documents |
 | --max-async | 4 | Maximum async operations |
 | --max-tokens | 32768 | Maximum token size |
 | --embedding-dim | 1024 | Embedding dimensions |
 | --max-embed-tokens | 8192 | Maximum embedding token size |
-| --input-file | ./book.txt | Initial input file |
-| --log-level | INFO | Logging level |
-| --key | none | Access Key to protect the lightrag service |
+| --timeout | None | Timeout in seconds (useful when using slow AI). Use None for infinite timeout |
+| --log-level | INFO | Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL) |
+| --key | None | API key for authentication. Protects lightrag server against unauthorized access |
+| --ssl | False | Enable HTTPS |
+| --ssl-certfile | None | Path to SSL certificate file (required if --ssl is enabled) |
+| --ssl-keyfile | None | Path to SSL private key file (required if --ssl is enabled) |
 
-#### Ollama Server Options
 
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| --host | 0.0.0.0 | RAG server host |
-| --port | 9621 | RAG server port |
-| --model | mistral-nemo:latest | LLM model name |
-| --embedding-model | bge-m3:latest | Embedding model name |
-| --ollama-host | http://localhost:11434 | Ollama backend URL |
-| --working-dir | ./rag_storage | Working directory for RAG |
-| --max-async | 4 | Maximum async operations |
-| --max-tokens | 32768 | Maximum token size |
-| --embedding-dim | 1024 | Embedding dimensions |
-| --max-embed-tokens | 8192 | Maximum embedding token size |
-| --input-file | ./book.txt | Initial input file |
-| --log-level | INFO | Logging level |
-| --key | none | Access Key to protect the lightrag service |
 
-#### OpenAI Server Options
+For protecting the server using an authentication key, you can also use an environment variable named `LIGHTRAG_API_KEY`.
+### Example Usage
 
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| --host | 0.0.0.0 | RAG server host |
-| --port | 9621 | RAG server port |
-| --model | gpt-4 | OpenAI model name |
-| --embedding-model | text-embedding-3-large | OpenAI embedding model |
-| --working-dir | ./rag_storage | Working directory for RAG |
-| --max-tokens | 32768 | Maximum token size |
-| --max-embed-tokens | 8192 | Maximum embedding token size |
-| --input-dir | ./inputs | Input directory for documents |
-| --log-level | INFO | Logging level |
-| --key | none | Access Key to protect the lightrag service |
+#### Running a Lightrag server with ollama default local server as llm and embedding backends
 
-#### OpenAI AZURE Server Options
+Ollama is the default backend for both llm and embedding, so by default you can run lightrag-server with no parameters and the default ones will be used. Make sure ollama is installed and is running and default models are already installed on ollama.
 
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| --host | 0.0.0.0 | Server host |
-| --port | 9621 | Server port |
-| --model | gpt-4 | OpenAI model name |
-| --embedding-model | text-embedding-3-large | OpenAI embedding model |
-| --working-dir | ./rag_storage | Working directory for RAG |
-| --max-tokens | 32768 | Maximum token size |
-| --max-embed-tokens | 8192 | Maximum embedding token size |
-| --input-dir | ./inputs | Input directory for documents |
-| --enable-cache | True | Enable response cache |
-| --log-level | INFO | Logging level |
-| --key | none | Access Key to protect the lightrag service |
+```bash
+# Run lightrag with ollama, mistral-nemo:latest for llm, and bge-m3:latest for embedding
+lightrag-server
 
+# Using specific models (ensure they are installed in your ollama instance)
+lightrag-server --llm-model adrienbrault/nous-hermes2theta-llama3-8b:f16 --embedding-model nomic-embed-text --embedding-dim 1024
 
-For protecting the server using an authentication key, you can also use an environment variable named `LIGHTRAG_API_KEY`.
-### Example Usage
+# Using an authentication key
+lightrag-server --key my-key
 
-#### LoLLMs RAG Server
+# Using lollms for llm and ollama for embedding
+lightrag-server --llm-binding lollms
+```
+
+#### Running a Lightrag server with lollms default local server as llm and embedding backends
 
 ```bash
-# Custom configuration with specific model and working directory
-lollms-lightrag-server --model mistral-nemo --port 8080 --working-dir ./custom_rag
+# Run lightrag with lollms, mistral-nemo:latest for llm, and bge-m3:latest for embedding, use lollms for both llm and embedding
+lightrag-server --llm-binding lollms --embedding-binding lollms
 
-# Using specific models (ensure they are installed in your LoLLMs instance)
-lollms-lightrag-server --model mistral-nemo:latest --embedding-model bge-m3 --embedding-dim 1024
+# Using specific models (ensure they are installed in your ollama instance)
+lightrag-server --llm-binding lollms --llm-model adrienbrault/nous-hermes2theta-llama3-8b:f16 --embedding-binding lollms --embedding-model nomic-embed-text --embedding-dim 1024
 
-# Using specific models and an authentication key
-lollms-lightrag-server --model mistral-nemo:latest --embedding-model bge-m3 --embedding-dim 1024 --key ky-mykey
+# Using an authentication key
+lightrag-server --key my-key
 
+# Using lollms for llm and openai for embedding
+lightrag-server --llm-binding lollms --embedding-binding openai --embedding-model text-embedding-3-small
 ```
 
-#### Ollama RAG Server
+
+#### Running a Lightrag server with openai server as llm and embedding backends
 
 ```bash
-# Custom configuration with specific model and working directory
-ollama-lightrag-server --model mistral-nemo:latest --port 8080 --working-dir ./custom_rag
+# Run lightrag with lollms, GPT-4o-mini  for llm, and text-embedding-3-small for embedding, use openai for both llm and embedding
+lightrag-server --llm-binding openai --llm-model GPT-4o-mini --embedding-binding openai --embedding-model text-embedding-3-small
+
+# Using an authentication key
+lightrag-server --llm-binding openai --llm-model GPT-4o-mini --embedding-binding openai --embedding-model text-embedding-3-small --key my-key
 
-# Using specific models (ensure they are installed in your Ollama instance)
-ollama-lightrag-server --model mistral-nemo:latest --embedding-model bge-m3 --embedding-dim 1024
+# Using lollms for llm and openai for embedding
+lightrag-server --llm-binding lollms --embedding-binding openai --embedding-model text-embedding-3-small
 ```
 
-#### OpenAI RAG Server
+#### Running a Lightrag server with azure openai server as llm and embedding backends
 
 ```bash
-# Using GPT-4 with text-embedding-3-large
-openai-lightrag-server --port 9624 --model gpt-4 --embedding-model text-embedding-3-large
-```
-#### Azure OpenAI RAG Server
-```bash
-# Using GPT-4 with text-embedding-3-large
-azure-openai-lightrag-server --model gpt-4o --port 8080 --working-dir ./custom_rag --embedding-model text-embedding-3-large
-```
+# Run lightrag with lollms, GPT-4o-mini  for llm, and text-embedding-3-small for embedding, use openai for both llm and embedding
+lightrag-server --llm-binding azure_openai --llm-model GPT-4o-mini --embedding-binding openai --embedding-model text-embedding-3-small
 
+# Using an authentication key
+lightrag-server --llm-binding azure_openai --llm-model GPT-4o-mini --embedding-binding azure_openai --embedding-model text-embedding-3-small --key my-key
+
+# Using lollms for llm and azure_openai for embedding
+lightrag-server --llm-binding lollms --embedding-binding azure_openai --embedding-model text-embedding-3-small
+```
 
 **Important Notes:**
 - For LoLLMs: Make sure the specified models are installed in your LoLLMs instance
@@ -1069,10 +1067,7 @@ azure-openai-lightrag-server --model gpt-4o --port 8080 --working-dir ./custom_r
 
 For help on any server, use the --help flag:
 ```bash
-lollms-lightrag-server --help
-ollama-lightrag-server --help
-openai-lightrag-server --help
-azure-openai-lightrag-server --help
+lightrag-server --help
 ```
 
 Note: If you don't need the API functionality, you can install the base package without API support using:
@@ -1092,7 +1087,7 @@ Query the RAG system with options for different search modes.
 ```bash
 curl -X POST "http://localhost:9621/query" \
     -H "Content-Type: application/json" \
-    -d '{"query": "Your question here", "mode": "hybrid"}'
+    -d '{"query": "Your question here", "mode": "hybrid", ""}'
 ```
 
 #### POST /query/stream

examples/test_split_by_character.ipynb

@@ -0,0 +1,1296 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "4b5690db12e34685",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:40:58.307102Z",
+     "start_time": "2025-01-09T03:40:51.935233Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import logging\n",
+    "import numpy as np\n",
+    "from lightrag import LightRAG, QueryParam\n",
+    "from lightrag.llm import openai_complete_if_cache, openai_embedding\n",
+    "from lightrag.utils import EmbeddingFunc\n",
+    "import nest_asyncio"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dd17956ec322b361",
+   "metadata": {},
+   "source": "#### split by character"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "8c8ee7c061bf9159",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:41:13.961167Z",
+     "start_time": "2025-01-09T03:41:13.958357Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "nest_asyncio.apply()\n",
+    "WORKING_DIR = \"../../llm_rag/paper_db/R000088_test1\"\n",
+    "logging.basicConfig(format=\"%(levelname)s:%(message)s\", level=logging.INFO)\n",
+    "if not os.path.exists(WORKING_DIR):\n",
+    "    os.mkdir(WORKING_DIR)\n",
+    "API = os.environ.get(\"DOUBAO_API_KEY\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "a5009d16e0851dca",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:41:16.862036Z",
+     "start_time": "2025-01-09T03:41:16.859306Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "async def llm_model_func(\n",
+    "    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs\n",
+    ") -> str:\n",
+    "    return await openai_complete_if_cache(\n",
+    "        \"ep-20241218114828-2tlww\",\n",
+    "        prompt,\n",
+    "        system_prompt=system_prompt,\n",
+    "        history_messages=history_messages,\n",
+    "        api_key=API,\n",
+    "        base_url=\"https://ark.cn-beijing.volces.com/api/v3\",\n",
+    "        **kwargs,\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "async def embedding_func(texts: list[str]) -> np.ndarray:\n",
+    "    return await openai_embedding(\n",
+    "        texts,\n",
+    "        model=\"ep-20241231173413-pgjmk\",\n",
+    "        api_key=API,\n",
+    "        base_url=\"https://ark.cn-beijing.volces.com/api/v3\",\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "397fcad24ce4d0ed",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:41:24.950307Z",
+     "start_time": "2025-01-09T03:41:24.940353Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:lightrag:Logger initialized for working directory: ../../llm_rag/paper_db/R000088_test1\n",
+      "INFO:lightrag:Load KV llm_response_cache with 0 data\n",
+      "INFO:lightrag:Load KV full_docs with 0 data\n",
+      "INFO:lightrag:Load KV text_chunks with 0 data\n",
+      "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test1/vdb_entities.json'} 0 data\n",
+      "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test1/vdb_relationships.json'} 0 data\n",
+      "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test1/vdb_chunks.json'} 0 data\n",
+      "INFO:lightrag:Loaded document status storage with 0 records\n"
+     ]
+    }
+   ],
+   "source": [
+    "rag = LightRAG(\n",
+    "    working_dir=WORKING_DIR,\n",
+    "    llm_model_func=llm_model_func,\n",
+    "    embedding_func=EmbeddingFunc(\n",
+    "        embedding_dim=4096, max_token_size=8192, func=embedding_func\n",
+    "    ),\n",
+    "    chunk_token_size=512,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "1dc3603677f7484d",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:41:37.947456Z",
+     "start_time": "2025-01-09T03:41:37.941901Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "with open(\n",
+    "    \"../../llm_rag/example/R000088/auto/R000088_full_txt.md\", \"r\", encoding=\"utf-8\"\n",
+    ") as f:\n",
+    "    content = f.read()\n",
+    "\n",
+    "\n",
+    "async def embedding_func(texts: list[str]) -> np.ndarray:\n",
+    "    return await openai_embedding(\n",
+    "        texts,\n",
+    "        model=\"ep-20241231173413-pgjmk\",\n",
+    "        api_key=API,\n",
+    "        base_url=\"https://ark.cn-beijing.volces.com/api/v3\",\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "async def get_embedding_dim():\n",
+    "    test_text = [\"This is a test sentence.\"]\n",
+    "    embedding = await embedding_func(test_text)\n",
+    "    embedding_dim = embedding.shape[1]\n",
+    "    return embedding_dim"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "6844202606acfbe5",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:41:39.608541Z",
+     "start_time": "2025-01-09T03:41:39.165057Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n"
+     ]
+    }
+   ],
+   "source": [
+    "embedding_dimension = await get_embedding_dim()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "d6273839d9681403",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:44:34.295345Z",
+     "start_time": "2025-01-09T03:41:48.324171Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:lightrag:Processing 1 new unique documents\n",
+      "Processing batch 1:   0%|          | 0/1 [00:00<?, ?it/s]INFO:lightrag:Inserting 35 vectors to chunks\n",
+      "\n",
+      "Generating embeddings:   0%|          | 0/2 [00:00<?, ?batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  50%|█████     | 1/2 [00:00<00:00,  1.36batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings: 100%|██████████| 2/2 [00:04<00:00,  2.25s/batch]\u001b[A\n",
+      "\n",
+      "Extracting entities from chunks:   0%|          | 0/35 [00:00<?, ?chunk/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠙ Processed 1 chunks, 1 entities(duplicated), 0 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:   3%|▎         | 1/35 [00:04<02:47,  4.93s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠹ Processed 2 chunks, 2 entities(duplicated), 0 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:   6%|▌         | 2/35 [00:05<01:18,  2.37s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠸ Processed 3 chunks, 9 entities(duplicated), 5 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:   9%|▊         | 3/35 [00:26<05:43, 10.73s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠼ Processed 4 chunks, 16 entities(duplicated), 11 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  11%|█▏        | 4/35 [00:26<03:24,  6.60s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠴ Processed 5 chunks, 24 entities(duplicated), 18 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  14%|█▍        | 5/35 [00:33<03:24,  6.82s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠦ Processed 6 chunks, 35 entities(duplicated), 28 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  17%|█▋        | 6/35 [00:42<03:38,  7.53s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠧ Processed 7 chunks, 47 entities(duplicated), 36 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  20%|██        | 7/35 [00:43<02:28,  5.31s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠇ Processed 8 chunks, 61 entities(duplicated), 49 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  23%|██▎       | 8/35 [00:45<01:52,  4.16s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠏ Processed 9 chunks, 81 entities(duplicated), 49 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠋ Processed 10 chunks, 90 entities(duplicated), 62 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  29%|██▊       | 10/35 [00:46<01:06,  2.64s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠙ Processed 11 chunks, 101 entities(duplicated), 80 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  31%|███▏      | 11/35 [00:52<01:19,  3.31s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠹ Processed 12 chunks, 108 entities(duplicated), 85 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  34%|███▍      | 12/35 [00:54<01:11,  3.12s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠸ Processed 13 chunks, 120 entities(duplicated), 100 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  37%|███▋      | 13/35 [00:59<01:18,  3.55s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠼ Processed 14 chunks, 131 entities(duplicated), 110 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  40%|████      | 14/35 [01:00<00:59,  2.82s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠴ Processed 15 chunks, 143 entities(duplicated), 110 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  43%|████▎     | 15/35 [01:02<00:52,  2.64s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠦ Processed 16 chunks, 162 entities(duplicated), 124 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  46%|████▌     | 16/35 [01:05<00:53,  2.80s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠧ Processed 17 chunks, 174 entities(duplicated), 132 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  49%|████▊     | 17/35 [01:06<00:39,  2.19s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠇ Processed 18 chunks, 185 entities(duplicated), 137 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  51%|█████▏    | 18/35 [01:12<00:53,  3.15s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠏ Processed 19 chunks, 193 entities(duplicated), 149 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  54%|█████▍    | 19/35 [01:18<01:06,  4.14s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠋ Processed 20 chunks, 205 entities(duplicated), 158 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  57%|█████▋    | 20/35 [01:19<00:50,  3.35s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠙ Processed 21 chunks, 220 entities(duplicated), 187 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  60%|██████    | 21/35 [01:27<01:02,  4.47s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠹ Processed 22 chunks, 247 entities(duplicated), 216 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  63%|██████▎   | 22/35 [01:30<00:54,  4.16s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠸ Processed 23 chunks, 260 entities(duplicated), 230 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  66%|██████▌   | 23/35 [01:34<00:48,  4.05s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠼ Processed 24 chunks, 291 entities(duplicated), 253 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  69%|██████▊   | 24/35 [01:38<00:44,  4.03s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠴ Processed 25 chunks, 304 entities(duplicated), 262 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  71%|███████▏  | 25/35 [01:41<00:36,  3.67s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠦ Processed 26 chunks, 313 entities(duplicated), 271 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  74%|███████▍  | 26/35 [01:41<00:24,  2.76s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠧ Processed 27 chunks, 321 entities(duplicated), 283 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  77%|███████▋  | 27/35 [01:47<00:28,  3.52s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠇ Processed 28 chunks, 333 entities(duplicated), 290 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  80%|████████  | 28/35 [01:52<00:28,  4.08s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠏ Processed 29 chunks, 348 entities(duplicated), 307 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  83%|████████▎ | 29/35 [01:59<00:29,  4.88s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠋ Processed 30 chunks, 362 entities(duplicated), 329 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  86%|████████▌ | 30/35 [02:02<00:21,  4.29s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠙ Processed 31 chunks, 373 entities(duplicated), 337 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  89%|████████▊ | 31/35 [02:03<00:13,  3.28s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠹ Processed 32 chunks, 390 entities(duplicated), 369 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  91%|█████████▏| 32/35 [02:03<00:07,  2.55s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠸ Processed 33 chunks, 405 entities(duplicated), 378 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  94%|█████████▍| 33/35 [02:07<00:05,  2.84s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠼ Processed 34 chunks, 435 entities(duplicated), 395 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  97%|█████████▋| 34/35 [02:10<00:02,  2.94s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠴ Processed 35 chunks, 456 entities(duplicated), 440 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks: 100%|██████████| 35/35 [02:23<00:00,  4.10s/chunk]\u001b[A\n",
+      "INFO:lightrag:Inserting entities into storage...\n",
+      "\n",
+      "Inserting entities: 100%|██████████| 324/324 [00:00<00:00, 17456.96entity/s]\n",
+      "INFO:lightrag:Inserting relationships into storage...\n",
+      "\n",
+      "Inserting relationships: 100%|██████████| 427/427 [00:00<00:00, 29956.31relationship/s]\n",
+      "INFO:lightrag:Inserting 324 vectors to entities\n",
+      "\n",
+      "Generating embeddings:   0%|          | 0/11 [00:00<?, ?batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:   9%|▉         | 1/11 [00:00<00:06,  1.48batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  18%|█▊        | 2/11 [00:02<00:11,  1.25s/batch]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  27%|██▋       | 3/11 [00:02<00:06,  1.17batch/s]\u001b[A\n",
+      "Generating embeddings:  36%|███▋      | 4/11 [00:03<00:04,  1.50batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  45%|████▌     | 5/11 [00:03<00:03,  1.78batch/s]\u001b[A\n",
+      "Generating embeddings:  55%|█████▍    | 6/11 [00:03<00:02,  2.01batch/s]\u001b[A\n",
+      "Generating embeddings:  64%|██████▎   | 7/11 [00:04<00:01,  2.19batch/s]\u001b[A\n",
+      "Generating embeddings:  73%|███████▎  | 8/11 [00:04<00:01,  2.31batch/s]\u001b[A\n",
+      "Generating embeddings:  82%|████████▏ | 9/11 [00:04<00:00,  2.41batch/s]\u001b[A\n",
+      "Generating embeddings:  91%|█████████ | 10/11 [00:05<00:00,  2.48batch/s]\u001b[A\n",
+      "Generating embeddings: 100%|██████████| 11/11 [00:05<00:00,  1.91batch/s]\u001b[A\n",
+      "INFO:lightrag:Inserting 427 vectors to relationships\n",
+      "\n",
+      "Generating embeddings:   0%|          | 0/14 [00:00<?, ?batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:   7%|▋         | 1/14 [00:01<00:14,  1.11s/batch]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  14%|█▍        | 2/14 [00:02<00:14,  1.18s/batch]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  21%|██▏       | 3/14 [00:02<00:08,  1.23batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  29%|██▊       | 4/14 [00:03<00:06,  1.56batch/s]\u001b[A\n",
+      "Generating embeddings:  36%|███▌      | 5/14 [00:03<00:04,  1.85batch/s]\u001b[A\n",
+      "Generating embeddings:  43%|████▎     | 6/14 [00:03<00:03,  2.05batch/s]\u001b[A\n",
+      "Generating embeddings:  50%|█████     | 7/14 [00:04<00:03,  2.23batch/s]\u001b[A\n",
+      "Generating embeddings:  57%|█████▋    | 8/14 [00:04<00:02,  2.37batch/s]\u001b[A\n",
+      "Generating embeddings:  64%|██████▍   | 9/14 [00:04<00:02,  2.46batch/s]\u001b[A\n",
+      "Generating embeddings:  71%|███████▏  | 10/14 [00:05<00:01,  2.54batch/s]\u001b[A\n",
+      "Generating embeddings:  79%|███████▊  | 11/14 [00:05<00:01,  2.59batch/s]\u001b[A\n",
+      "Generating embeddings:  86%|████████▌ | 12/14 [00:06<00:00,  2.64batch/s]\u001b[A\n",
+      "Generating embeddings:  93%|█████████▎| 13/14 [00:06<00:00,  2.65batch/s]\u001b[A\n",
+      "Generating embeddings: 100%|██████████| 14/14 [00:06<00:00,  2.05batch/s]\u001b[A\n",
+      "INFO:lightrag:Writing graph with 333 nodes, 427 edges\n",
+      "Processing batch 1: 100%|██████████| 1/1 [02:45<00:00, 165.90s/it]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# rag.insert(content)\n",
+    "rag.insert(content, split_by_character=\"\\n#\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "c4f9ae517151a01d",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:45:11.668987Z",
+     "start_time": "2025-01-09T03:45:11.664744Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "prompt1 = \"\"\"你是一名经验丰富的论文分析科学家，你的任务是对一篇英文学术研究论文进行关键信息提取并深入分析。\n",
+    "请按照以下步骤进行分析：\n",
+    "1. 该文献主要研究的问题是什么？\n",
+    "2. 该文献采用什么方法进行分析？\n",
+    "3. 该文献的主要结论是什么？\n",
+    "首先在<分析>标签中，针对每个问题详��分析你的思考过程。然后在<回答>标签中给出所有问题的最终答案。\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "7a6491385b050095",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:45:40.829111Z",
+     "start_time": "2025-01-09T03:45:13.530298Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:lightrag:Local query uses 5 entites, 12 relations, 3 text units\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:lightrag:Global query uses 8 entites, 5 relations, 4 text units\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "<分析>\n",
+      "1. **该文献主要研究的问题是什么？**\n",
+      "    - 思考过程：通过浏览论文内容，查找作者明确阐述研究目的的部分。文中多处提及“Our study was performed to explore whether folic acid treatment was associated with cancer outcomes and all-cause mortality after extended follow-up”，表明作者旨在探究叶酸治疗与癌症结局及全因死亡率之间的关系，尤其是在经过长期随访后。\n",
+      "2. **该文献采用什么方法进行分析？**\n",
+      "    - 思考过程：寻找描述研究方法和数据分析过程的段落。文中提到“Survival curves were constructed using the Kaplan-Meier method and differences in survival between groups were analyzed using the log-rank test. Estimates of hazard ratios (HRs) with 95% CIs were obtained by using Cox proportional hazards regression models stratified by trial”，可以看出作者使用了Kaplan-Meier法构建生存曲线、log-rank检验分析组间生存差异以及Cox比例风险回归模型估计风险比等方法。\n",
+      "3. **该文献的主要结论是什么？**\n",
+      "    - 思考过程：定位到论文中总结结论的部分，如“Conclusion Treatment with folic acid plus vitamin $\\mathsf{B}_{12}$ was associated with increased cancer outcomes and all-cause mortality in patients with ischemic heart disease in Norway, where there is no folic acid fortification of foods”，可知作者得出叶酸加维生素$\\mathsf{B}_{12}$治疗与癌症结局和全因死亡率增加有关的结论。\n",
+      "<回答>\n",
+      "1. 该文献主要研究的问题是：叶酸治疗与癌症结局及全因死亡率之间的关系，尤其是在经过长期随访后，叶酸治疗是否与癌症结局和全因死亡率相关。\n",
+      "2. 该文献采用的分析方法包括：使用Kaplan-Meier法构建生存曲线、log-rank检验分析组间生存差异、Cox比例风险回归模型估计风险比等。\n",
+      "3. 该文献的主要结论是：在挪威没有叶酸强化食品的情况下，叶酸加维生素$\\mathsf{B}_{12}$治疗与缺血性心脏病患者的癌症结局和全因死亡率增加有关。\n",
+      "\n",
+      "**参考文献**\n",
+      "- [VD] In2Norwegianhomocysteine-lowering trialsamongpatientswithischemicheart disease, there was a statistically nonsignificantincreaseincancerincidenceinthe groupsassignedtofolicacidtreatment.15,16 Our study was performed to explore whetherfolicacidtreatmentwasassociatedwithcanceroutcomesandall-cause mortality after extended follow-up.\n",
+      "- [VD] Survivalcurveswereconstructedusing theKaplan-Meiermethodanddifferences insurvivalbetweengroupswereanalyzed usingthelog-ranktest.Estimatesofhazard ratios (HRs) with $95\\%$ CIs were obtainedbyusingCoxproportionalhazards regressionmodelsstratifiedbytrial.\n",
+      "- [VD] Conclusion Treatment with folic acid plus vitamin $\\mathsf{B}_{12}$ was associated with increased cancer outcomes and all-cause mortality in patients with ischemic heart disease in Norway, where there is no folic acid fortification of foods.\n"
+     ]
+    }
+   ],
+   "source": [
+    "resp = rag.query(prompt1, param=QueryParam(mode=\"mix\", top_k=5))\n",
+    "print(resp)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4e5bfad24cb721a8",
+   "metadata": {},
+   "source": "#### split by character only"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "44e2992dc95f8ce0",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:47:40.988796Z",
+     "start_time": "2025-01-09T03:47:40.982648Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "WORKING_DIR = \"../../llm_rag/paper_db/R000088_test2\"\n",
+    "if not os.path.exists(WORKING_DIR):\n",
+    "    os.mkdir(WORKING_DIR)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "62c63385d2d973d5",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:51:39.951329Z",
+     "start_time": "2025-01-09T03:49:15.218976Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:lightrag:Logger initialized for working directory: ../../llm_rag/paper_db/R000088_test2\n",
+      "INFO:lightrag:Load KV llm_response_cache with 0 data\n",
+      "INFO:lightrag:Load KV full_docs with 0 data\n",
+      "INFO:lightrag:Load KV text_chunks with 0 data\n",
+      "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test2/vdb_entities.json'} 0 data\n",
+      "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test2/vdb_relationships.json'} 0 data\n",
+      "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test2/vdb_chunks.json'} 0 data\n",
+      "INFO:lightrag:Loaded document status storage with 0 records\n",
+      "INFO:lightrag:Processing 1 new unique documents\n",
+      "Processing batch 1:   0%|          | 0/1 [00:00<?, ?it/s]INFO:lightrag:Inserting 12 vectors to chunks\n",
+      "\n",
+      "Generating embeddings:   0%|          | 0/1 [00:00<?, ?batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings: 100%|██████████| 1/1 [00:02<00:00,  2.95s/batch]\u001b[A\n",
+      "\n",
+      "Extracting entities from chunks:   0%|          | 0/12 [00:00<?, ?chunk/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠙ Processed 1 chunks, 0 entities(duplicated), 0 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:   8%|▊         | 1/12 [00:03<00:43,  3.93s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠹ Processed 2 chunks, 8 entities(duplicated), 8 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  17%|█▋        | 2/12 [00:29<02:44, 16.46s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠸ Processed 3 chunks, 17 entities(duplicated), 15 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  25%|██▌       | 3/12 [00:30<01:25,  9.45s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠼ Processed 4 chunks, 27 entities(duplicated), 22 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  33%|███▎      | 4/12 [00:39<01:16,  9.52s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠴ Processed 5 chunks, 36 entities(duplicated), 33 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  42%|████▏     | 5/12 [00:40<00:43,  6.24s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠦ Processed 6 chunks, 49 entities(duplicated), 42 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  50%|█████     | 6/12 [00:49<00:43,  7.33s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠧ Processed 7 chunks, 62 entities(duplicated), 65 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  58%|█████▊    | 7/12 [01:05<00:50, 10.05s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠇ Processed 8 chunks, 81 entities(duplicated), 90 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  67%|██████▋   | 8/12 [01:23<00:50, 12.69s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠏ Processed 9 chunks, 99 entities(duplicated), 117 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  75%|███████▌  | 9/12 [01:32<00:34, 11.54s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠋ Processed 10 chunks, 123 entities(duplicated), 140 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  83%|████████▎ | 10/12 [01:48<00:25, 12.79s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠙ Processed 11 chunks, 158 entities(duplicated), 174 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks:  92%|█████████▏| 11/12 [02:03<00:13, 13.50s/chunk]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "⠹ Processed 12 chunks, 194 entities(duplicated), 221 relations(duplicated)\r"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Extracting entities from chunks: 100%|██████████| 12/12 [02:13<00:00, 11.15s/chunk]\u001b[A\n",
+      "INFO:lightrag:Inserting entities into storage...\n",
+      "\n",
+      "Inserting entities: 100%|██████████| 170/170 [00:00<00:00, 11610.25entity/s]\n",
+      "INFO:lightrag:Inserting relationships into storage...\n",
+      "\n",
+      "Inserting relationships: 100%|██████████| 218/218 [00:00<00:00, 15913.51relationship/s]\n",
+      "INFO:lightrag:Inserting 170 vectors to entities\n",
+      "\n",
+      "Generating embeddings:   0%|          | 0/6 [00:00<?, ?batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  17%|█▋        | 1/6 [00:01<00:05,  1.10s/batch]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  33%|███▎      | 2/6 [00:02<00:04,  1.07s/batch]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  50%|█████     | 3/6 [00:02<00:02,  1.33batch/s]\u001b[A\n",
+      "Generating embeddings:  67%|██████▋   | 4/6 [00:02<00:01,  1.67batch/s]\u001b[A\n",
+      "Generating embeddings:  83%|████████▎ | 5/6 [00:03<00:00,  1.95batch/s]\u001b[A\n",
+      "Generating embeddings: 100%|██████████| 6/6 [00:03<00:00,  1.66batch/s]\u001b[A\n",
+      "INFO:lightrag:Inserting 218 vectors to relationships\n",
+      "\n",
+      "Generating embeddings:   0%|          | 0/7 [00:00<?, ?batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  14%|█▍        | 1/7 [00:01<00:10,  1.74s/batch]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  29%|██▊       | 2/7 [00:02<00:05,  1.04s/batch]\u001b[A\n",
+      "Generating embeddings:  43%|████▎     | 3/7 [00:02<00:02,  1.35batch/s]\u001b[A\n",
+      "Generating embeddings:  57%|█████▋    | 4/7 [00:03<00:01,  1.69batch/s]\u001b[AINFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "\n",
+      "Generating embeddings:  71%|███████▏  | 5/7 [00:03<00:01,  1.96batch/s]\u001b[A\n",
+      "Generating embeddings:  86%|████████▌ | 6/7 [00:03<00:00,  2.17batch/s]\u001b[A\n",
+      "Generating embeddings: 100%|██████████| 7/7 [00:04<00:00,  1.68batch/s]\u001b[A\n",
+      "INFO:lightrag:Writing graph with 174 nodes, 218 edges\n",
+      "Processing batch 1: 100%|██████████| 1/1 [02:24<00:00, 144.69s/it]\n"
+     ]
+    }
+   ],
+   "source": [
+    "rag = LightRAG(\n",
+    "    working_dir=WORKING_DIR,\n",
+    "    llm_model_func=llm_model_func,\n",
+    "    embedding_func=EmbeddingFunc(\n",
+    "        embedding_dim=4096, max_token_size=8192, func=embedding_func\n",
+    "    ),\n",
+    "    chunk_token_size=512,\n",
+    ")\n",
+    "\n",
+    "# rag.insert(content)\n",
+    "rag.insert(content, split_by_character=\"\\n#\", split_by_character_only=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "3c7aa9836d8d43c7",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-01-09T03:52:37.000418Z",
+     "start_time": "2025-01-09T03:52:09.933584Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:lightrag:Local query uses 5 entites, 3 relations, 2 text units\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n",
+      "INFO:lightrag:Global query uses 9 entites, 5 relations, 4 text units\n",
+      "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "<分析>\n",
+      "- **该文献主要研究的问题是什么？**\n",
+      "    - **思考过程**：通过浏览论文的标题、摘要、引言等部分，寻找关于研究目的和问题的描述。论文标题为“Cancer Incidence and Mortality After Treatment With Folic Acid and Vitamin B12”，摘要中的“Objective”部分明确指出研究目的是“To evaluate effects of treatment with B vitamins on cancer outcomes and all-cause mortality in 2 randomized controlled trials”。因此，可以确定该文献主要研究的问题是评估B族维生素治疗对两项随机对照试验中癌症结局和全因死亡率的影响。\n",
+      "- **该文献采用什么方法进行分析？**\n",
+      "    - **思考过程**：在论文的“METHODS”部分详细描述了研究方法。文中提到这是一个对两项随机、双盲、安慰剂对照临床试验（Norwegian Vitamin [NORVIT] trial和Western Norway B Vitamin Intervention Trial [WENBIT]）数据的联合分析，并进行了观察性的试验后随访。具体包括对参与者进行分组干预（不同剂量的叶酸、维生素B12、维生素B6或安慰剂），收集临床信息和血样，分析循环B族维生素、同型半胱氨酸和可替宁等指标，并进行基因分型等，还涉及到多种统计分析方法，如计算预期癌症发生率、构建生存曲线、进行Cox比例风险回归模型分析等。\n",
+      "- **该文献的主要结论是什么？**\n",
+      "    - **思考过程**：在论文的“Results”和“Conclusion”部分寻找主要结论。研究结果表明，在治疗期间，接受叶酸加维生素B12治疗的参与者血清叶酸浓度显著增加，且在后续随访中，该组癌症发病率、癌症死亡率和全因死亡率均有所上升，主要是肺癌发病率增加，而维生素B6治疗未显示出显著影响。结论部分明确指出“Treatment with folic acid plus vitamin $\\mathsf{B}_{12}$ was associated with increased cancer outcomes and all-cause mortality in patients with ischemic heart disease in Norway, where there is no folic acid fortification of foods”。\n",
+      "</分析>\n",
+      "\n",
+      "<回答>\n",
+      "- **主要研究问题**：评估B族维生素治疗对两项随机对照试验中癌症结局和全因死亡率的影响。\n",
+      "- **研究方法**：采用对两项随机、双盲、安慰剂对照临床试验（Norwegian Vitamin [NORVIT] trial和Western Norway B Vitamin Intervention Trial [WENBIT]）数据的联合分析，并进行观察性的试验后随访，涉及分组干预、多种指标检测以及多种统计分析方法。\n",
+      "- **主要结论**：在挪威（食品中未添加叶酸），对于缺血性心脏病患者，叶酸加维生素B12治疗与癌症结局和全因死亡率的增加有关，而维生素B6治疗未显示出显著影响。\n",
+      "\n",
+      "**参考文献**\n",
+      "- [VD] Cancer Incidence and Mortality After Treatment With Folic Acid and Vitamin B12\n",
+      "- [VD] METHODS Study Design, Participants, and Study Intervention\n",
+      "- [VD] RESULTS\n",
+      "- [VD] Conclusion\n",
+      "- [VD] Objective To evaluate effects of treatment with B vitamins on cancer outcomes and all-cause mortality in 2 randomized controlled trials.\n"
+     ]
+    }
+   ],
+   "source": [
+    "resp = rag.query(prompt1, param=QueryParam(mode=\"mix\", top_k=5))\n",
+    "print(resp)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7ba6fa79a2550d10",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

lightrag/__init__.py

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

lightrag/api/azure_openai_lightrag_server.py

@@ -1,532 +0,0 @@
-from fastapi import FastAPI, HTTPException, File, UploadFile, Form
-from pydantic import BaseModel
-import asyncio
-import logging
-import argparse
-from lightrag import LightRAG, QueryParam
-from lightrag.llm import (
-    azure_openai_complete_if_cache,
-    azure_openai_embedding,
-)
-from lightrag.utils import EmbeddingFunc
-from typing import Optional, List
-from enum import Enum
-from pathlib import Path
-import shutil
-import aiofiles
-from ascii_colors import trace_exception
-import os
-from dotenv import load_dotenv
-import inspect
-import json
-from fastapi.responses import StreamingResponse
-
-from fastapi import Depends, Security
-from fastapi.security import APIKeyHeader
-from fastapi.middleware.cors import CORSMiddleware
-
-from starlette.status import HTTP_403_FORBIDDEN
-
-load_dotenv()
-
-AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION")
-AZURE_OPENAI_DEPLOYMENT = os.getenv("AZURE_OPENAI_DEPLOYMENT")
-AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
-AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
-
-AZURE_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_EMBEDDING_DEPLOYMENT")
-AZURE_EMBEDDING_API_VERSION = os.getenv("AZURE_EMBEDDING_API_VERSION")
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="LightRAG FastAPI Server with OpenAI integration"
-    )
-
-    # Server configuration
-    parser.add_argument(
-        "--host", default="0.0.0.0", help="Server host (default: 0.0.0.0)"
-    )
-    parser.add_argument(
-        "--port", type=int, default=9621, help="Server port (default: 9621)"
-    )
-
-    # Directory configuration
-    parser.add_argument(
-        "--working-dir",
-        default="./rag_storage",
-        help="Working directory for RAG storage (default: ./rag_storage)",
-    )
-    parser.add_argument(
-        "--input-dir",
-        default="./inputs",
-        help="Directory containing input documents (default: ./inputs)",
-    )
-
-    # Model configuration
-    parser.add_argument(
-        "--model", default="gpt-4o", help="OpenAI model name (default: gpt-4o)"
-    )
-    parser.add_argument(
-        "--embedding-model",
-        default="text-embedding-3-large",
-        help="OpenAI embedding model (default: text-embedding-3-large)",
-    )
-
-    # RAG configuration
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=32768,
-        help="Maximum token size (default: 32768)",
-    )
-    parser.add_argument(
-        "--max-embed-tokens",
-        type=int,
-        default=8192,
-        help="Maximum embedding token size (default: 8192)",
-    )
-    parser.add_argument(
-        "--enable-cache",
-        default=True,
-        help="Enable response cache (default: True)",
-    )
-    # Logging configuration
-    parser.add_argument(
-        "--log-level",
-        default="INFO",
-        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
-        help="Logging level (default: INFO)",
-    )
-
-    parser.add_argument(
-        "--key",
-        type=str,
-        help="API key for authentication. This protects lightrag server against unauthorized access",
-        default=None,
-    )
-
-    return parser.parse_args()
-
-
-class DocumentManager:
-    """Handles document operations and tracking"""
-
-    def __init__(self, input_dir: str, supported_extensions: tuple = (".txt", ".md")):
-        self.input_dir = Path(input_dir)
-        self.supported_extensions = supported_extensions
-        self.indexed_files = set()
-
-        # Create input directory if it doesn't exist
-        self.input_dir.mkdir(parents=True, exist_ok=True)
-
-    def scan_directory(self) -> List[Path]:
-        """Scan input directory for new files"""
-        new_files = []
-        for ext in self.supported_extensions:
-            for file_path in self.input_dir.rglob(f"*{ext}"):
-                if file_path not in self.indexed_files:
-                    new_files.append(file_path)
-        return new_files
-
-    def mark_as_indexed(self, file_path: Path):
-        """Mark a file as indexed"""
-        self.indexed_files.add(file_path)
-
-    def is_supported_file(self, filename: str) -> bool:
-        """Check if file type is supported"""
-        return any(filename.lower().endswith(ext) for ext in self.supported_extensions)
-
-
-# Pydantic models
-class SearchMode(str, Enum):
-    naive = "naive"
-    local = "local"
-    global_ = "global"
-    hybrid = "hybrid"
-
-
-class QueryRequest(BaseModel):
-    query: str
-    mode: SearchMode = SearchMode.hybrid
-    only_need_context: bool = False
-    # stream: bool = False
-
-
-class QueryResponse(BaseModel):
-    response: str
-
-
-class InsertTextRequest(BaseModel):
-    text: str
-    description: Optional[str] = None
-
-
-class InsertResponse(BaseModel):
-    status: str
-    message: str
-    document_count: int
-
-
-def get_api_key_dependency(api_key: Optional[str]):
-    if not api_key:
-        # If no API key is configured, return a dummy dependency that always succeeds
-        async def no_auth():
-            return None
-
-        return no_auth
-
-    # If API key is configured, use proper authentication
-    api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
-
-    async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
-        if not api_key_header_value:
-            raise HTTPException(
-                status_code=HTTP_403_FORBIDDEN, detail="API Key required"
-            )
-        if api_key_header_value != api_key:
-            raise HTTPException(
-                status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
-            )
-        return api_key_header_value
-
-    return api_key_auth
-
-
-async def get_embedding_dim(embedding_model: str) -> int:
-    """Get embedding dimensions for the specified model"""
-    test_text = ["This is a test sentence."]
-    embedding = await azure_openai_embedding(test_text, model=embedding_model)
-    return embedding.shape[1]
-
-
-def create_app(args):
-    # Setup logging
-    logging.basicConfig(
-        format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
-    )
-
-    # Check if API key is provided either through env var or args
-    api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
-
-    # Initialize FastAPI
-    app = FastAPI(
-        title="LightRAG API",
-        description="API for querying text using LightRAG with separate storage and input directories"
-        + "(With authentication)"
-        if api_key
-        else "",
-        version="1.0.0",
-        openapi_tags=[{"name": "api"}],
-    )
-
-    # Add CORS middleware
-    app.add_middleware(
-        CORSMiddleware,
-        allow_origins=["*"],
-        allow_credentials=True,
-        allow_methods=["*"],
-        allow_headers=["*"],
-    )
-
-    # Create the optional API key dependency
-    optional_api_key = get_api_key_dependency(api_key)
-
-    # Create working directory if it doesn't exist
-    Path(args.working_dir).mkdir(parents=True, exist_ok=True)
-
-    # Initialize document manager
-    doc_manager = DocumentManager(args.input_dir)
-
-    # Get embedding dimensions
-    embedding_dim = asyncio.run(get_embedding_dim(args.embedding_model))
-
-    async def async_openai_complete(
-        prompt, system_prompt=None, history_messages=[], **kwargs
-    ):
-        """Async wrapper for OpenAI completion"""
-        kwargs.pop("keyword_extraction", None)
-
-        return await azure_openai_complete_if_cache(
-            args.model,
-            prompt,
-            system_prompt=system_prompt,
-            history_messages=history_messages,
-            base_url=AZURE_OPENAI_ENDPOINT,
-            api_key=AZURE_OPENAI_API_KEY,
-            api_version=AZURE_OPENAI_API_VERSION,
-            **kwargs,
-        )
-
-    # Initialize RAG with OpenAI configuration
-    rag = LightRAG(
-        enable_llm_cache=args.enable_cache,
-        working_dir=args.working_dir,
-        llm_model_func=async_openai_complete,
-        llm_model_name=args.model,
-        llm_model_max_token_size=args.max_tokens,
-        embedding_func=EmbeddingFunc(
-            embedding_dim=embedding_dim,
-            max_token_size=args.max_embed_tokens,
-            func=lambda texts: azure_openai_embedding(
-                texts, model=args.embedding_model
-            ),
-        ),
-    )
-
-    @app.on_event("startup")
-    async def startup_event():
-        """Index all files in input directory during startup"""
-        try:
-            new_files = doc_manager.scan_directory()
-            for file_path in new_files:
-                try:
-                    # Use async file reading
-                    async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
-                        content = await f.read()
-                        # Use the async version of insert directly
-                        await rag.ainsert(content)
-                        doc_manager.mark_as_indexed(file_path)
-                        logging.info(f"Indexed file: {file_path}")
-                except Exception as e:
-                    trace_exception(e)
-                    logging.error(f"Error indexing file {file_path}: {str(e)}")
-
-            logging.info(f"Indexed {len(new_files)} documents from {args.input_dir}")
-
-        except Exception as e:
-            logging.error(f"Error during startup indexing: {str(e)}")
-
-    @app.post("/documents/scan", dependencies=[Depends(optional_api_key)])
-    async def scan_for_new_documents():
-        """Manually trigger scanning for new documents"""
-        try:
-            new_files = doc_manager.scan_directory()
-            indexed_count = 0
-
-            for file_path in new_files:
-                try:
-                    with open(file_path, "r", encoding="utf-8") as f:
-                        content = f.read()
-                        await rag.ainsert(content)
-                        doc_manager.mark_as_indexed(file_path)
-                        indexed_count += 1
-                except Exception as e:
-                    logging.error(f"Error indexing file {file_path}: {str(e)}")
-
-            return {
-                "status": "success",
-                "indexed_count": indexed_count,
-                "total_documents": len(doc_manager.indexed_files),
-            }
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post("/resetcache", dependencies=[Depends(optional_api_key)])
-    async def reset_cache():
-        """Manually reset cache"""
-        try:
-            cachefile = args.working_dir + "/kv_store_llm_response_cache.json"
-            if os.path.exists(cachefile):
-                with open(cachefile, "w") as f:
-                    f.write("{}")
-            return {"status": "success"}
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post("/documents/upload", dependencies=[Depends(optional_api_key)])
-    async def upload_to_input_dir(file: UploadFile = File(...)):
-        """Upload a file to the input directory"""
-        try:
-            if not doc_manager.is_supported_file(file.filename):
-                raise HTTPException(
-                    status_code=400,
-                    detail=f"Unsupported file type. Supported types: {doc_manager.supported_extensions}",
-                )
-
-            file_path = doc_manager.input_dir / file.filename
-            with open(file_path, "wb") as buffer:
-                shutil.copyfileobj(file.file, buffer)
-
-            # Immediately index the uploaded file
-            with open(file_path, "r", encoding="utf-8") as f:
-                content = f.read()
-                await rag.ainsert(content)
-                doc_manager.mark_as_indexed(file_path)
-
-            return {
-                "status": "success",
-                "message": f"File uploaded and indexed: {file.filename}",
-                "total_documents": len(doc_manager.indexed_files),
-            }
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
-    )
-    async def query_text(request: QueryRequest):
-        try:
-            response = await rag.aquery(
-                request.query,
-                param=QueryParam(
-                    mode=request.mode,
-                    stream=False,
-                    only_need_context=request.only_need_context,
-                ),
-            )
-            return QueryResponse(response=response)
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post("/query/stream", dependencies=[Depends(optional_api_key)])
-    async def query_text_stream(request: QueryRequest):
-        try:
-            response = await rag.aquery(
-                request.query,
-                param=QueryParam(
-                    mode=request.mode,
-                    stream=True,
-                    only_need_context=request.only_need_context,
-                ),
-            )
-            if inspect.isasyncgen(response):
-
-                async def stream_generator():
-                    async for chunk in response:
-                        yield json.dumps({"data": chunk}) + "\n"
-
-                return StreamingResponse(
-                    stream_generator(), media_type="application/json"
-                )
-            else:
-                return QueryResponse(response=response)
-
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/documents/text",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def insert_text(request: InsertTextRequest):
-        try:
-            await rag.ainsert(request.text)
-            return InsertResponse(
-                status="success",
-                message="Text successfully inserted",
-                document_count=1,
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/documents/file",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
-        try:
-            content = await file.read()
-
-            if file.filename.endswith((".txt", ".md")):
-                text = content.decode("utf-8")
-                rag.insert(text)
-            else:
-                raise HTTPException(
-                    status_code=400,
-                    detail="Unsupported file type. Only .txt and .md files are supported",
-                )
-
-            return InsertResponse(
-                status="success",
-                message=f"File '{file.filename}' successfully inserted",
-                document_count=1,
-            )
-        except UnicodeDecodeError:
-            raise HTTPException(status_code=400, detail="File encoding not supported")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/documents/batch",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def insert_batch(files: List[UploadFile] = File(...)):
-        try:
-            inserted_count = 0
-            failed_files = []
-
-            for file in files:
-                try:
-                    content = await file.read()
-                    if file.filename.endswith((".txt", ".md")):
-                        text = content.decode("utf-8")
-                        rag.insert(text)
-                        inserted_count += 1
-                    else:
-                        failed_files.append(f"{file.filename} (unsupported type)")
-                except Exception as e:
-                    failed_files.append(f"{file.filename} ({str(e)})")
-
-            status_message = f"Successfully inserted {inserted_count} documents"
-            if failed_files:
-                status_message += f". Failed files: {', '.join(failed_files)}"
-
-            return InsertResponse(
-                status="success" if inserted_count > 0 else "partial_success",
-                message=status_message,
-                document_count=len(files),
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.delete(
-        "/documents",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def clear_documents():
-        try:
-            rag.text_chunks = []
-            rag.entities_vdb = None
-            rag.relationships_vdb = None
-            return InsertResponse(
-                status="success",
-                message="All documents cleared successfully",
-                document_count=0,
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.get("/health", dependencies=[Depends(optional_api_key)])
-    async def get_status():
-        """Get current system status"""
-        return {
-            "status": "healthy",
-            "working_directory": str(args.working_dir),
-            "input_directory": str(args.input_dir),
-            "indexed_files": len(doc_manager.indexed_files),
-            "configuration": {
-                "model": args.model,
-                "embedding_model": args.embedding_model,
-                "max_tokens": args.max_tokens,
-                "embedding_dim": embedding_dim,
-            },
-        }
-
-    return app
-
-
-def main():
-    args = parse_args()
-    import uvicorn
-
-    app = create_app(args)
-    uvicorn.run(app, host=args.host, port=args.port)
-
-
-if __name__ == "__main__":
-    main()

lightrag/api/{lollms_lightrag_server.py → lightrag_server.py}

@@ -4,8 +4,12 @@ import logging
 import argparse
 from lightrag import LightRAG, QueryParam
 from lightrag.llm import lollms_model_complete, lollms_embed
+from lightrag.llm import ollama_model_complete, ollama_embed
+from lightrag.llm import openai_complete_if_cache, openai_embedding
+from lightrag.llm import azure_openai_complete_if_cache, azure_openai_embedding
+
 from lightrag.utils import EmbeddingFunc
-from typing import Optional, List
+from typing import Optional, List, Union
 from enum import Enum
 from pathlib import Path
 import shutil
@@ -18,6 +22,19 @@ from fastapi.security import APIKeyHeader
 from fastapi.middleware.cors import CORSMiddleware
 
 from starlette.status import HTTP_403_FORBIDDEN
+import pipmaster as pm
+
+
+def get_default_host(binding_type: str) -> str:
+    default_hosts = {
+        "ollama": "http://localhost:11434",
+        "lollms": "http://localhost:9600",
+        "azure_openai": "https://api.openai.com/v1",
+        "openai": "https://api.openai.com/v1",
+    }
+    return default_hosts.get(
+        binding_type, "http://localhost:11434"
+    )  # fallback to ollama if unknown
 
 
 def parse_args():
@@ -25,6 +42,22 @@ def parse_args():
         description="LightRAG FastAPI Server with separate working and input directories"
     )
 
+    # Start by the bindings
+    parser.add_argument(
+        "--llm-binding",
+        default="ollama",
+        help="LLM binding to be used. Supported: lollms, ollama, openai (default: ollama)",
+    )
+    parser.add_argument(
+        "--embedding-binding",
+        default="ollama",
+        help="Embedding binding to be used. Supported: lollms, ollama, openai (default: ollama)",
+    )
+
+    # Parse just these arguments first
+    temp_args, _ = parser.parse_known_args()
+
+    # Add remaining arguments with dynamic defaults for hosts
     # Server configuration
     parser.add_argument(
         "--host", default="0.0.0.0", help="Server host (default: 0.0.0.0)"
@@ -45,23 +78,45 @@ def parse_args():
         help="Directory containing input documents (default: ./inputs)",
     )
 
-    # Model configuration
+    # LLM Model configuration
+    default_llm_host = get_default_host(temp_args.llm_binding)
+    parser.add_argument(
+        "--llm-binding-host",
+        default=default_llm_host,
+        help=f"llm server host URL (default: {default_llm_host})",
+    )
+
     parser.add_argument(
-        "--model",
+        "--llm-model",
         default="mistral-nemo:latest",
         help="LLM model name (default: mistral-nemo:latest)",
     )
+
+    # Embedding model configuration
+    default_embedding_host = get_default_host(temp_args.embedding_binding)
+    parser.add_argument(
+        "--embedding-binding-host",
+        default=default_embedding_host,
+        help=f"embedding server host URL (default: {default_embedding_host})",
+    )
+
     parser.add_argument(
         "--embedding-model",
         default="bge-m3:latest",
         help="Embedding model name (default: bge-m3:latest)",
     )
+
+    def timeout_type(value):
+        if value is None or value == "None":
+            return None
+        return int(value)
+
     parser.add_argument(
-        "--lollms-host",
-        default="http://localhost:9600",
-        help="lollms host URL (default: http://localhost:9600)",
+        "--timeout",
+        default=None,
+        type=timeout_type,
+        help="Timeout in seconds (useful when using slow AI). Use None for infinite timeout",
     )
-
     # RAG configuration
     parser.add_argument(
         "--max-async", type=int, default=4, help="Maximum async operations (default: 4)"
@@ -100,13 +155,31 @@ def parse_args():
         default=None,
     )
 
+    # Optional https parameters
+    parser.add_argument(
+        "--ssl", action="store_true", help="Enable HTTPS (default: False)"
+    )
+    parser.add_argument(
+        "--ssl-certfile",
+        default=None,
+        help="Path to SSL certificate file (required if --ssl is enabled)",
+    )
+    parser.add_argument(
+        "--ssl-keyfile",
+        default=None,
+        help="Path to SSL private key file (required if --ssl is enabled)",
+    )
     return parser.parse_args()
 
 
 class DocumentManager:
     """Handles document operations and tracking"""
 
-    def __init__(self, input_dir: str, supported_extensions: tuple = (".txt", ".md")):
+    def __init__(
+        self,
+        input_dir: str,
+        supported_extensions: tuple = (".txt", ".md", ".pdf", ".docx", ".pptx"),
+    ):
         self.input_dir = Path(input_dir)
         self.supported_extensions = supported_extensions
         self.indexed_files = set()
@@ -188,6 +261,24 @@ def get_api_key_dependency(api_key: Optional[str]):
 
 
 def create_app(args):
+    # Verify that bindings arer correctly setup
+    if args.llm_binding not in ["lollms", "ollama", "openai"]:
+        raise Exception("llm binding not supported")
+
+    if args.embedding_binding not in ["lollms", "ollama", "openai"]:
+        raise Exception("embedding binding not supported")
+
+    # Add SSL validation
+    if args.ssl:
+        if not args.ssl_certfile or not args.ssl_keyfile:
+            raise Exception(
+                "SSL certificate and key files must be provided when SSL is enabled"
+            )
+        if not os.path.exists(args.ssl_certfile):
+            raise Exception(f"SSL certificate file not found: {args.ssl_certfile}")
+        if not os.path.exists(args.ssl_keyfile):
+            raise Exception(f"SSL key file not found: {args.ssl_keyfile}")
+
     # Setup logging
     logging.basicConfig(
         format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
@@ -203,7 +294,7 @@ def create_app(args):
         + "(With authentication)"
         if api_key
         else "",
-        version="1.0.0",
+        version="1.0.2",
         openapi_tags=[{"name": "api"}],
     )
 
@@ -228,23 +319,122 @@ def create_app(args):
     # Initialize RAG
     rag = LightRAG(
         working_dir=args.working_dir,
-        llm_model_func=lollms_model_complete,
-        llm_model_name=args.model,
+        llm_model_func=lollms_model_complete
+        if args.llm_binding == "lollms"
+        else ollama_model_complete
+        if args.llm_binding == "ollama"
+        else azure_openai_complete_if_cache
+        if args.llm_binding == "azure_openai"
+        else openai_complete_if_cache,
+        llm_model_name=args.llm_model,
         llm_model_max_async=args.max_async,
         llm_model_max_token_size=args.max_tokens,
         llm_model_kwargs={
-            "host": args.lollms_host,
+            "host": args.llm_binding_host,
+            "timeout": args.timeout,
             "options": {"num_ctx": args.max_tokens},
         },
         embedding_func=EmbeddingFunc(
             embedding_dim=args.embedding_dim,
             max_token_size=args.max_embed_tokens,
             func=lambda texts: lollms_embed(
-                texts, embed_model=args.embedding_model, host=args.lollms_host
+                texts,
+                embed_model=args.embedding_model,
+                host=args.embedding_binding_host,
+            )
+            if args.llm_binding == "lollms"
+            else ollama_embed(
+                texts,
+                embed_model=args.embedding_model,
+                host=args.embedding_binding_host,
+            )
+            if args.llm_binding == "ollama"
+            else azure_openai_embedding(
+                texts,
+                model=args.embedding_model,  # no host is used for openai
+            )
+            if args.llm_binding == "azure_openai"
+            else openai_embedding(
+                texts,
+                model=args.embedding_model,  # no host is used for openai
             ),
         ),
     )
 
+    async def index_file(file_path: Union[str, Path]) -> None:
+        """Index all files inside the folder with support for multiple file formats
+
+        Args:
+            file_path: Path to the file to be indexed (str or Path object)
+
+        Raises:
+            ValueError: If file format is not supported
+            FileNotFoundError: If file doesn't exist
+        """
+        if not pm.is_installed("aiofiles"):
+            pm.install("aiofiles")
+
+        # Convert to Path object if string
+        file_path = Path(file_path)
+
+        # Check if file exists
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        content = ""
+        # Get file extension in lowercase
+        ext = file_path.suffix.lower()
+
+        match ext:
+            case ".txt" | ".md":
+                # Text files handling
+                async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
+                    content = await f.read()
+
+            case ".pdf":
+                if not pm.is_installed("pypdf2"):
+                    pm.install("pypdf2")
+                from pypdf2 import PdfReader
+
+                # PDF handling
+                reader = PdfReader(str(file_path))
+                content = ""
+                for page in reader.pages:
+                    content += page.extract_text() + "\n"
+
+            case ".docx":
+                if not pm.is_installed("docx"):
+                    pm.install("docx")
+                from docx import Document
+
+                # Word document handling
+                doc = Document(file_path)
+                content = "\n".join([paragraph.text for paragraph in doc.paragraphs])
+
+            case ".pptx":
+                if not pm.is_installed("pptx"):
+                    pm.install("pptx")
+                from pptx import Presentation
+
+                # PowerPoint handling
+                prs = Presentation(file_path)
+                content = ""
+                for slide in prs.slides:
+                    for shape in slide.shapes:
+                        if hasattr(shape, "text"):
+                            content += shape.text + "\n"
+
+            case _:
+                raise ValueError(f"Unsupported file format: {ext}")
+
+        # Insert content into RAG system
+        if content:
+            await rag.ainsert(content)
+            doc_manager.mark_as_indexed(file_path)
+            logging.info(f"Successfully indexed file: {file_path}")
+        else:
+            logging.warning(f"No content extracted from file: {file_path}")
+
     @app.on_event("startup")
     async def startup_event():
         """Index all files in input directory during startup"""
@@ -252,13 +442,7 @@ def create_app(args):
             new_files = doc_manager.scan_directory()
             for file_path in new_files:
                 try:
-                    # Use async file reading
-                    async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
-                        content = await f.read()
-                        # Use the async version of insert directly
-                        await rag.ainsert(content)
-                        doc_manager.mark_as_indexed(file_path)
-                        logging.info(f"Indexed file: {file_path}")
+                    await index_file(file_path)
                 except Exception as e:
                     trace_exception(e)
                     logging.error(f"Error indexing file {file_path}: {str(e)}")
@@ -277,11 +461,8 @@ def create_app(args):
 
             for file_path in new_files:
                 try:
-                    with open(file_path, "r", encoding="utf-8") as f:
-                        content = f.read()
-                        await rag.ainsert(content)
-                        doc_manager.mark_as_indexed(file_path)
-                        indexed_count += 1
+                    await index_file(file_path)
+                    indexed_count += 1
                 except Exception as e:
                     logging.error(f"Error indexing file {file_path}: {str(e)}")
 
@@ -308,10 +489,7 @@ def create_app(args):
                 shutil.copyfileobj(file.file, buffer)
 
             # Immediately index the uploaded file
-            with open(file_path, "r", encoding="utf-8") as f:
-                content = f.read()
-                await rag.ainsert(content)
-                doc_manager.mark_as_indexed(file_path)
+            await index_file(file_path)
 
             return {
                 "status": "success",
@@ -372,11 +550,11 @@ def create_app(args):
     )
     async def insert_text(request: InsertTextRequest):
         try:
-            rag.insert(request.text)
+            await rag.ainsert(request.text)
             return InsertResponse(
                 status="success",
                 message="Text successfully inserted",
-                document_count=len(rag),
+                document_count=1,
             )
         except Exception as e:
             raise HTTPException(status_code=500, detail=str(e))
@@ -387,26 +565,103 @@ def create_app(args):
         dependencies=[Depends(optional_api_key)],
     )
     async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
+        """Insert a file directly into the RAG system
+
+        Args:
+            file: Uploaded file
+            description: Optional description of the file
+
+        Returns:
+            InsertResponse: Status of the insertion operation
+
+        Raises:
+            HTTPException: For unsupported file types or processing errors
+        """
         try:
-            content = await file.read()
+            content = ""
+            # Get file extension in lowercase
+            ext = Path(file.filename).suffix.lower()
+
+            match ext:
+                case ".txt" | ".md":
+                    # Text files handling
+                    text_content = await file.read()
+                    content = text_content.decode("utf-8")
+
+                case ".pdf":
+                    if not pm.is_installed("pypdf2"):
+                        pm.install("pypdf2")
+                    from pypdf2 import PdfReader
+                    from io import BytesIO
+
+                    # Read PDF from memory
+                    pdf_content = await file.read()
+                    pdf_file = BytesIO(pdf_content)
+                    reader = PdfReader(pdf_file)
+                    content = ""
+                    for page in reader.pages:
+                        content += page.extract_text() + "\n"
+
+                case ".docx":
+                    if not pm.is_installed("docx"):
+                        pm.install("docx")
+                    from docx import Document
+                    from io import BytesIO
+
+                    # Read DOCX from memory
+                    docx_content = await file.read()
+                    docx_file = BytesIO(docx_content)
+                    doc = Document(docx_file)
+                    content = "\n".join(
+                        [paragraph.text for paragraph in doc.paragraphs]
+                    )
+
+                case ".pptx":
+                    if not pm.is_installed("pptx"):
+                        pm.install("pptx")
+                    from pptx import Presentation
+                    from io import BytesIO
+
+                    # Read PPTX from memory
+                    pptx_content = await file.read()
+                    pptx_file = BytesIO(pptx_content)
+                    prs = Presentation(pptx_file)
+                    content = ""
+                    for slide in prs.slides:
+                        for shape in slide.shapes:
+                            if hasattr(shape, "text"):
+                                content += shape.text + "\n"
+
+                case _:
+                    raise HTTPException(
+                        status_code=400,
+                        detail=f"Unsupported file type. Supported types: {doc_manager.supported_extensions}",
+                    )
+
+            # Insert content into RAG system
+            if content:
+                # Add description if provided
+                if description:
+                    content = f"{description}\n\n{content}"
+
+                await rag.ainsert(content)
+                logging.info(f"Successfully indexed file: {file.filename}")
 
-            if file.filename.endswith((".txt", ".md")):
-                text = content.decode("utf-8")
-                await rag.ainsert(text)
+                return InsertResponse(
+                    status="success",
+                    message=f"File '{file.filename}' successfully inserted",
+                    document_count=1,
+                )
             else:
                 raise HTTPException(
                     status_code=400,
-                    detail="Unsupported file type. Only .txt and .md files are supported",
+                    detail="No content could be extracted from the file",
                 )
 
-            return InsertResponse(
-                status="success",
-                message=f"File '{file.filename}' successfully inserted",
-                document_count=1,
-            )
         except UnicodeDecodeError:
             raise HTTPException(status_code=400, detail="File encoding not supported")
         except Exception as e:
+            logging.error(f"Error processing file {file.filename}: {str(e)}")
             raise HTTPException(status_code=500, detail=str(e))
 
     @app.post(
@@ -415,32 +670,110 @@ def create_app(args):
         dependencies=[Depends(optional_api_key)],
     )
     async def insert_batch(files: List[UploadFile] = File(...)):
+        """Process multiple files in batch mode
+
+        Args:
+            files: List of files to process
+
+        Returns:
+            InsertResponse: Status of the batch insertion operation
+
+        Raises:
+            HTTPException: For processing errors
+        """
         try:
             inserted_count = 0
             failed_files = []
 
             for file in files:
                 try:
-                    content = await file.read()
-                    if file.filename.endswith((".txt", ".md")):
-                        text = content.decode("utf-8")
-                        await rag.ainsert(text)
+                    content = ""
+                    ext = Path(file.filename).suffix.lower()
+
+                    match ext:
+                        case ".txt" | ".md":
+                            text_content = await file.read()
+                            content = text_content.decode("utf-8")
+
+                        case ".pdf":
+                            if not pm.is_installed("pypdf2"):
+                                pm.install("pypdf2")
+                            from pypdf2 import PdfReader
+                            from io import BytesIO
+
+                            pdf_content = await file.read()
+                            pdf_file = BytesIO(pdf_content)
+                            reader = PdfReader(pdf_file)
+                            for page in reader.pages:
+                                content += page.extract_text() + "\n"
+
+                        case ".docx":
+                            if not pm.is_installed("docx"):
+                                pm.install("docx")
+                            from docx import Document
+                            from io import BytesIO
+
+                            docx_content = await file.read()
+                            docx_file = BytesIO(docx_content)
+                            doc = Document(docx_file)
+                            content = "\n".join(
+                                [paragraph.text for paragraph in doc.paragraphs]
+                            )
+
+                        case ".pptx":
+                            if not pm.is_installed("pptx"):
+                                pm.install("pptx")
+                            from pptx import Presentation
+                            from io import BytesIO
+
+                            pptx_content = await file.read()
+                            pptx_file = BytesIO(pptx_content)
+                            prs = Presentation(pptx_file)
+                            for slide in prs.slides:
+                                for shape in slide.shapes:
+                                    if hasattr(shape, "text"):
+                                        content += shape.text + "\n"
+
+                        case _:
+                            failed_files.append(f"{file.filename} (unsupported type)")
+                            continue
+
+                    if content:
+                        await rag.ainsert(content)
                         inserted_count += 1
+                        logging.info(f"Successfully indexed file: {file.filename}")
                     else:
-                        failed_files.append(f"{file.filename} (unsupported type)")
+                        failed_files.append(f"{file.filename} (no content extracted)")
+
+                except UnicodeDecodeError:
+                    failed_files.append(f"{file.filename} (encoding error)")
                 except Exception as e:
                     failed_files.append(f"{file.filename} ({str(e)})")
-
-            status_message = f"Successfully inserted {inserted_count} documents"
-            if failed_files:
-                status_message += f". Failed files: {', '.join(failed_files)}"
+                    logging.error(f"Error processing file {file.filename}: {str(e)}")
+
+            # Prepare status message
+            if inserted_count == len(files):
+                status = "success"
+                status_message = f"Successfully inserted all {inserted_count} documents"
+            elif inserted_count > 0:
+                status = "partial_success"
+                status_message = f"Successfully inserted {inserted_count} out of {len(files)} documents"
+                if failed_files:
+                    status_message += f". Failed files: {', '.join(failed_files)}"
+            else:
+                status = "failure"
+                status_message = "No documents were successfully inserted"
+                if failed_files:
+                    status_message += f". Failed files: {', '.join(failed_files)}"
 
             return InsertResponse(
-                status="success" if inserted_count > 0 else "partial_success",
+                status=status,
                 message=status_message,
-                document_count=len(files),
+                document_count=inserted_count,
             )
+
         except Exception as e:
+            logging.error(f"Batch processing error: {str(e)}")
             raise HTTPException(status_code=500, detail=str(e))
 
     @app.delete(
@@ -470,10 +803,15 @@ def create_app(args):
             "input_directory": str(args.input_dir),
             "indexed_files": len(doc_manager.indexed_files),
             "configuration": {
-                "model": args.model,
+                # LLM configuration binding/host address (if applicable)/model (if applicable)
+                "llm_binding": args.llm_binding,
+                "llm_binding_host": args.llm_binding_host,
+                "llm_model": args.llm_model,
+                # embedding model configuration binding/host address (if applicable)/model (if applicable)
+                "embedding_binding": args.embedding_binding,
+                "embedding_binding_host": args.embedding_binding_host,
                 "embedding_model": args.embedding_model,
                 "max_tokens": args.max_tokens,
-                "lollms_host": args.lollms_host,
             },
         }
 
@@ -485,7 +823,19 @@ def main():
     import uvicorn
 
     app = create_app(args)
-    uvicorn.run(app, host=args.host, port=args.port)
+    uvicorn_config = {
+        "app": app,
+        "host": args.host,
+        "port": args.port,
+    }
+    if args.ssl:
+        uvicorn_config.update(
+            {
+                "ssl_certfile": args.ssl_certfile,
+                "ssl_keyfile": args.ssl_keyfile,
+            }
+        )
+    uvicorn.run(**uvicorn_config)
 
 
 if __name__ == "__main__":

lightrag/api/ollama_lightrag_server.py

@@ -1,491 +0,0 @@
-from fastapi import FastAPI, HTTPException, File, UploadFile, Form
-from pydantic import BaseModel
-import logging
-import argparse
-from lightrag import LightRAG, QueryParam
-from lightrag.llm import ollama_model_complete, ollama_embed
-from lightrag.utils import EmbeddingFunc
-from typing import Optional, List
-from enum import Enum
-from pathlib import Path
-import shutil
-import aiofiles
-from ascii_colors import trace_exception
-import os
-
-from fastapi import Depends, Security
-from fastapi.security import APIKeyHeader
-from fastapi.middleware.cors import CORSMiddleware
-
-from starlette.status import HTTP_403_FORBIDDEN
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="LightRAG FastAPI Server with separate working and input directories"
-    )
-
-    # Server configuration
-    parser.add_argument(
-        "--host", default="0.0.0.0", help="Server host (default: 0.0.0.0)"
-    )
-    parser.add_argument(
-        "--port", type=int, default=9621, help="Server port (default: 9621)"
-    )
-
-    # Directory configuration
-    parser.add_argument(
-        "--working-dir",
-        default="./rag_storage",
-        help="Working directory for RAG storage (default: ./rag_storage)",
-    )
-    parser.add_argument(
-        "--input-dir",
-        default="./inputs",
-        help="Directory containing input documents (default: ./inputs)",
-    )
-
-    # Model configuration
-    parser.add_argument(
-        "--model",
-        default="mistral-nemo:latest",
-        help="LLM model name (default: mistral-nemo:latest)",
-    )
-    parser.add_argument(
-        "--embedding-model",
-        default="bge-m3:latest",
-        help="Embedding model name (default: bge-m3:latest)",
-    )
-    parser.add_argument(
-        "--ollama-host",
-        default="http://localhost:11434",
-        help="Ollama host URL (default: http://localhost:11434)",
-    )
-
-    # RAG configuration
-    parser.add_argument(
-        "--max-async", type=int, default=4, help="Maximum async operations (default: 4)"
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=32768,
-        help="Maximum token size (default: 32768)",
-    )
-    parser.add_argument(
-        "--embedding-dim",
-        type=int,
-        default=1024,
-        help="Embedding dimensions (default: 1024)",
-    )
-    parser.add_argument(
-        "--max-embed-tokens",
-        type=int,
-        default=8192,
-        help="Maximum embedding token size (default: 8192)",
-    )
-
-    # Logging configuration
-    parser.add_argument(
-        "--log-level",
-        default="INFO",
-        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
-        help="Logging level (default: INFO)",
-    )
-    parser.add_argument(
-        "--key",
-        type=str,
-        help="API key for authentication. This protects lightrag server against unauthorized access",
-        default=None,
-    )
-
-    return parser.parse_args()
-
-
-class DocumentManager:
-    """Handles document operations and tracking"""
-
-    def __init__(self, input_dir: str, supported_extensions: tuple = (".txt", ".md")):
-        self.input_dir = Path(input_dir)
-        self.supported_extensions = supported_extensions
-        self.indexed_files = set()
-
-        # Create input directory if it doesn't exist
-        self.input_dir.mkdir(parents=True, exist_ok=True)
-
-    def scan_directory(self) -> List[Path]:
-        """Scan input directory for new files"""
-        new_files = []
-        for ext in self.supported_extensions:
-            for file_path in self.input_dir.rglob(f"*{ext}"):
-                if file_path not in self.indexed_files:
-                    new_files.append(file_path)
-        return new_files
-
-    def mark_as_indexed(self, file_path: Path):
-        """Mark a file as indexed"""
-        self.indexed_files.add(file_path)
-
-    def is_supported_file(self, filename: str) -> bool:
-        """Check if file type is supported"""
-        return any(filename.lower().endswith(ext) for ext in self.supported_extensions)
-
-
-# Pydantic models
-class SearchMode(str, Enum):
-    naive = "naive"
-    local = "local"
-    global_ = "global"
-    hybrid = "hybrid"
-
-
-class QueryRequest(BaseModel):
-    query: str
-    mode: SearchMode = SearchMode.hybrid
-    stream: bool = False
-    only_need_context: bool = False
-
-
-class QueryResponse(BaseModel):
-    response: str
-
-
-class InsertTextRequest(BaseModel):
-    text: str
-    description: Optional[str] = None
-
-
-class InsertResponse(BaseModel):
-    status: str
-    message: str
-    document_count: int
-
-
-def get_api_key_dependency(api_key: Optional[str]):
-    if not api_key:
-        # If no API key is configured, return a dummy dependency that always succeeds
-        async def no_auth():
-            return None
-
-        return no_auth
-
-    # If API key is configured, use proper authentication
-    api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
-
-    async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
-        if not api_key_header_value:
-            raise HTTPException(
-                status_code=HTTP_403_FORBIDDEN, detail="API Key required"
-            )
-        if api_key_header_value != api_key:
-            raise HTTPException(
-                status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
-            )
-        return api_key_header_value
-
-    return api_key_auth
-
-
-def create_app(args):
-    # Setup logging
-    logging.basicConfig(
-        format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
-    )
-
-    # Check if API key is provided either through env var or args
-    api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
-
-    # Initialize FastAPI
-    app = FastAPI(
-        title="LightRAG API",
-        description="API for querying text using LightRAG with separate storage and input directories"
-        + "(With authentication)"
-        if api_key
-        else "",
-        version="1.0.0",
-        openapi_tags=[{"name": "api"}],
-    )
-
-    # Add CORS middleware
-    app.add_middleware(
-        CORSMiddleware,
-        allow_origins=["*"],
-        allow_credentials=True,
-        allow_methods=["*"],
-        allow_headers=["*"],
-    )
-
-    # Create the optional API key dependency
-    optional_api_key = get_api_key_dependency(api_key)
-
-    # Create working directory if it doesn't exist
-    Path(args.working_dir).mkdir(parents=True, exist_ok=True)
-
-    # Initialize document manager
-    doc_manager = DocumentManager(args.input_dir)
-
-    # Initialize RAG
-    rag = LightRAG(
-        working_dir=args.working_dir,
-        llm_model_func=ollama_model_complete,
-        llm_model_name=args.model,
-        llm_model_max_async=args.max_async,
-        llm_model_max_token_size=args.max_tokens,
-        llm_model_kwargs={
-            "host": args.ollama_host,
-            "options": {"num_ctx": args.max_tokens},
-        },
-        embedding_func=EmbeddingFunc(
-            embedding_dim=args.embedding_dim,
-            max_token_size=args.max_embed_tokens,
-            func=lambda texts: ollama_embed(
-                texts, embed_model=args.embedding_model, host=args.ollama_host
-            ),
-        ),
-    )
-
-    @app.on_event("startup")
-    async def startup_event():
-        """Index all files in input directory during startup"""
-        try:
-            new_files = doc_manager.scan_directory()
-            for file_path in new_files:
-                try:
-                    # Use async file reading
-                    async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
-                        content = await f.read()
-                        # Use the async version of insert directly
-                        await rag.ainsert(content)
-                        doc_manager.mark_as_indexed(file_path)
-                        logging.info(f"Indexed file: {file_path}")
-                except Exception as e:
-                    trace_exception(e)
-                    logging.error(f"Error indexing file {file_path}: {str(e)}")
-
-            logging.info(f"Indexed {len(new_files)} documents from {args.input_dir}")
-
-        except Exception as e:
-            logging.error(f"Error during startup indexing: {str(e)}")
-
-    @app.post("/documents/scan", dependencies=[Depends(optional_api_key)])
-    async def scan_for_new_documents():
-        """Manually trigger scanning for new documents"""
-        try:
-            new_files = doc_manager.scan_directory()
-            indexed_count = 0
-
-            for file_path in new_files:
-                try:
-                    with open(file_path, "r", encoding="utf-8") as f:
-                        content = f.read()
-                        await rag.ainsert(content)
-                        doc_manager.mark_as_indexed(file_path)
-                        indexed_count += 1
-                except Exception as e:
-                    logging.error(f"Error indexing file {file_path}: {str(e)}")
-
-            return {
-                "status": "success",
-                "indexed_count": indexed_count,
-                "total_documents": len(doc_manager.indexed_files),
-            }
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post("/documents/upload", dependencies=[Depends(optional_api_key)])
-    async def upload_to_input_dir(file: UploadFile = File(...)):
-        """Upload a file to the input directory"""
-        try:
-            if not doc_manager.is_supported_file(file.filename):
-                raise HTTPException(
-                    status_code=400,
-                    detail=f"Unsupported file type. Supported types: {doc_manager.supported_extensions}",
-                )
-
-            file_path = doc_manager.input_dir / file.filename
-            with open(file_path, "wb") as buffer:
-                shutil.copyfileobj(file.file, buffer)
-
-            # Immediately index the uploaded file
-            with open(file_path, "r", encoding="utf-8") as f:
-                content = f.read()
-                await rag.ainsert(content)
-                doc_manager.mark_as_indexed(file_path)
-
-            return {
-                "status": "success",
-                "message": f"File uploaded and indexed: {file.filename}",
-                "total_documents": len(doc_manager.indexed_files),
-            }
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
-    )
-    async def query_text(request: QueryRequest):
-        try:
-            response = await rag.aquery(
-                request.query,
-                param=QueryParam(
-                    mode=request.mode,
-                    stream=request.stream,
-                    only_need_context=request.only_need_context,
-                ),
-            )
-
-            if request.stream:
-                result = ""
-                async for chunk in response:
-                    result += chunk
-                return QueryResponse(response=result)
-            else:
-                return QueryResponse(response=response)
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post("/query/stream", dependencies=[Depends(optional_api_key)])
-    async def query_text_stream(request: QueryRequest):
-        try:
-            response = rag.query(
-                request.query,
-                param=QueryParam(
-                    mode=request.mode,
-                    stream=True,
-                    only_need_context=request.only_need_context,
-                ),
-            )
-
-            async def stream_generator():
-                async for chunk in response:
-                    yield chunk
-
-            return stream_generator()
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/documents/text",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def insert_text(request: InsertTextRequest):
-        try:
-            await rag.ainsert(request.text)
-            return InsertResponse(
-                status="success",
-                message="Text successfully inserted",
-                document_count=len(rag),
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/documents/file",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
-        try:
-            content = await file.read()
-
-            if file.filename.endswith((".txt", ".md")):
-                text = content.decode("utf-8")
-                await rag.ainsert(text)
-            else:
-                raise HTTPException(
-                    status_code=400,
-                    detail="Unsupported file type. Only .txt and .md files are supported",
-                )
-
-            return InsertResponse(
-                status="success",
-                message=f"File '{file.filename}' successfully inserted",
-                document_count=1,
-            )
-        except UnicodeDecodeError:
-            raise HTTPException(status_code=400, detail="File encoding not supported")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/documents/batch",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def insert_batch(files: List[UploadFile] = File(...)):
-        try:
-            inserted_count = 0
-            failed_files = []
-
-            for file in files:
-                try:
-                    content = await file.read()
-                    if file.filename.endswith((".txt", ".md")):
-                        text = content.decode("utf-8")
-                        await rag.ainsert(text)
-                        inserted_count += 1
-                    else:
-                        failed_files.append(f"{file.filename} (unsupported type)")
-                except Exception as e:
-                    failed_files.append(f"{file.filename} ({str(e)})")
-
-            status_message = f"Successfully inserted {inserted_count} documents"
-            if failed_files:
-                status_message += f". Failed files: {', '.join(failed_files)}"
-
-            return InsertResponse(
-                status="success" if inserted_count > 0 else "partial_success",
-                message=status_message,
-                document_count=len(files),
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.delete(
-        "/documents",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def clear_documents():
-        try:
-            rag.text_chunks = []
-            rag.entities_vdb = None
-            rag.relationships_vdb = None
-            return InsertResponse(
-                status="success",
-                message="All documents cleared successfully",
-                document_count=0,
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.get("/health", dependencies=[Depends(optional_api_key)])
-    async def get_status():
-        """Get current system status"""
-        return {
-            "status": "healthy",
-            "working_directory": str(args.working_dir),
-            "input_directory": str(args.input_dir),
-            "indexed_files": len(doc_manager.indexed_files),
-            "configuration": {
-                "model": args.model,
-                "embedding_model": args.embedding_model,
-                "max_tokens": args.max_tokens,
-                "ollama_host": args.ollama_host,
-            },
-        }
-
-    return app
-
-
-def main():
-    args = parse_args()
-    import uvicorn
-
-    app = create_app(args)
-    uvicorn.run(app, host=args.host, port=args.port)
-
-
-if __name__ == "__main__":
-    main()

lightrag/api/openai_lightrag_server.py

@@ -1,506 +0,0 @@
-from fastapi import FastAPI, HTTPException, File, UploadFile, Form
-from pydantic import BaseModel
-import asyncio
-import logging
-import argparse
-from lightrag import LightRAG, QueryParam
-from lightrag.llm import openai_complete_if_cache, openai_embedding
-from lightrag.utils import EmbeddingFunc
-from typing import Optional, List
-from enum import Enum
-from pathlib import Path
-import shutil
-import aiofiles
-from ascii_colors import trace_exception
-import nest_asyncio
-
-import os
-
-from fastapi import Depends, Security
-from fastapi.security import APIKeyHeader
-from fastapi.middleware.cors import CORSMiddleware
-
-from starlette.status import HTTP_403_FORBIDDEN
-
-# Apply nest_asyncio to solve event loop issues
-nest_asyncio.apply()
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="LightRAG FastAPI Server with OpenAI integration"
-    )
-
-    # Server configuration
-    parser.add_argument(
-        "--host", default="0.0.0.0", help="Server host (default: 0.0.0.0)"
-    )
-    parser.add_argument(
-        "--port", type=int, default=9621, help="Server port (default: 9621)"
-    )
-
-    # Directory configuration
-    parser.add_argument(
-        "--working-dir",
-        default="./rag_storage",
-        help="Working directory for RAG storage (default: ./rag_storage)",
-    )
-    parser.add_argument(
-        "--input-dir",
-        default="./inputs",
-        help="Directory containing input documents (default: ./inputs)",
-    )
-
-    # Model configuration
-    parser.add_argument(
-        "--model", default="gpt-4", help="OpenAI model name (default: gpt-4)"
-    )
-    parser.add_argument(
-        "--embedding-model",
-        default="text-embedding-3-large",
-        help="OpenAI embedding model (default: text-embedding-3-large)",
-    )
-
-    # RAG configuration
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=32768,
-        help="Maximum token size (default: 32768)",
-    )
-    parser.add_argument(
-        "--max-embed-tokens",
-        type=int,
-        default=8192,
-        help="Maximum embedding token size (default: 8192)",
-    )
-
-    # Logging configuration
-    parser.add_argument(
-        "--log-level",
-        default="INFO",
-        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
-        help="Logging level (default: INFO)",
-    )
-
-    parser.add_argument(
-        "--key",
-        type=str,
-        help="API key for authentication. This protects lightrag server against unauthorized access",
-        default=None,
-    )
-
-    return parser.parse_args()
-
-
-class DocumentManager:
-    """Handles document operations and tracking"""
-
-    def __init__(self, input_dir: str, supported_extensions: tuple = (".txt", ".md")):
-        self.input_dir = Path(input_dir)
-        self.supported_extensions = supported_extensions
-        self.indexed_files = set()
-
-        # Create input directory if it doesn't exist
-        self.input_dir.mkdir(parents=True, exist_ok=True)
-
-    def scan_directory(self) -> List[Path]:
-        """Scan input directory for new files"""
-        new_files = []
-        for ext in self.supported_extensions:
-            for file_path in self.input_dir.rglob(f"*{ext}"):
-                if file_path not in self.indexed_files:
-                    new_files.append(file_path)
-        return new_files
-
-    def mark_as_indexed(self, file_path: Path):
-        """Mark a file as indexed"""
-        self.indexed_files.add(file_path)
-
-    def is_supported_file(self, filename: str) -> bool:
-        """Check if file type is supported"""
-        return any(filename.lower().endswith(ext) for ext in self.supported_extensions)
-
-
-# Pydantic models
-class SearchMode(str, Enum):
-    naive = "naive"
-    local = "local"
-    global_ = "global"
-    hybrid = "hybrid"
-
-
-class QueryRequest(BaseModel):
-    query: str
-    mode: SearchMode = SearchMode.hybrid
-    stream: bool = False
-    only_need_context: bool = False
-
-
-class QueryResponse(BaseModel):
-    response: str
-
-
-class InsertTextRequest(BaseModel):
-    text: str
-    description: Optional[str] = None
-
-
-class InsertResponse(BaseModel):
-    status: str
-    message: str
-    document_count: int
-
-
-def get_api_key_dependency(api_key: Optional[str]):
-    if not api_key:
-        # If no API key is configured, return a dummy dependency that always succeeds
-        async def no_auth():
-            return None
-
-        return no_auth
-
-    # If API key is configured, use proper authentication
-    api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
-
-    async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
-        if not api_key_header_value:
-            raise HTTPException(
-                status_code=HTTP_403_FORBIDDEN, detail="API Key required"
-            )
-        if api_key_header_value != api_key:
-            raise HTTPException(
-                status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
-            )
-        return api_key_header_value
-
-    return api_key_auth
-
-
-async def get_embedding_dim(embedding_model: str) -> int:
-    """Get embedding dimensions for the specified model"""
-    test_text = ["This is a test sentence."]
-    embedding = await openai_embedding(test_text, model=embedding_model)
-    return embedding.shape[1]
-
-
-def create_app(args):
-    # Setup logging
-    logging.basicConfig(
-        format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
-    )
-
-    # Check if API key is provided either through env var or args
-    api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
-
-    # Initialize FastAPI
-    app = FastAPI(
-        title="LightRAG API",
-        description="API for querying text using LightRAG with separate storage and input directories"
-        + "(With authentication)"
-        if api_key
-        else "",
-        version="1.0.0",
-        openapi_tags=[{"name": "api"}],
-    )
-
-    # Add CORS middleware
-    app.add_middleware(
-        CORSMiddleware,
-        allow_origins=["*"],
-        allow_credentials=True,
-        allow_methods=["*"],
-        allow_headers=["*"],
-    )
-
-    # Create the optional API key dependency
-    optional_api_key = get_api_key_dependency(api_key)
-
-    # Add CORS middleware
-    app.add_middleware(
-        CORSMiddleware,
-        allow_origins=["*"],
-        allow_credentials=True,
-        allow_methods=["*"],
-        allow_headers=["*"],
-    )
-
-    # Create working directory if it doesn't exist
-    Path(args.working_dir).mkdir(parents=True, exist_ok=True)
-
-    # Initialize document manager
-    doc_manager = DocumentManager(args.input_dir)
-
-    # Get embedding dimensions
-    embedding_dim = asyncio.run(get_embedding_dim(args.embedding_model))
-
-    async def async_openai_complete(
-        prompt, system_prompt=None, history_messages=[], **kwargs
-    ):
-        """Async wrapper for OpenAI completion"""
-        return await openai_complete_if_cache(
-            args.model,
-            prompt,
-            system_prompt=system_prompt,
-            history_messages=history_messages,
-            **kwargs,
-        )
-
-    # Initialize RAG with OpenAI configuration
-    rag = LightRAG(
-        working_dir=args.working_dir,
-        llm_model_func=async_openai_complete,
-        llm_model_name=args.model,
-        llm_model_max_token_size=args.max_tokens,
-        embedding_func=EmbeddingFunc(
-            embedding_dim=embedding_dim,
-            max_token_size=args.max_embed_tokens,
-            func=lambda texts: openai_embedding(texts, model=args.embedding_model),
-        ),
-    )
-
-    @app.on_event("startup")
-    async def startup_event():
-        """Index all files in input directory during startup"""
-        try:
-            new_files = doc_manager.scan_directory()
-            for file_path in new_files:
-                try:
-                    # Use async file reading
-                    async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
-                        content = await f.read()
-                        # Use the async version of insert directly
-                        await rag.ainsert(content)
-                        doc_manager.mark_as_indexed(file_path)
-                        logging.info(f"Indexed file: {file_path}")
-                except Exception as e:
-                    trace_exception(e)
-                    logging.error(f"Error indexing file {file_path}: {str(e)}")
-
-            logging.info(f"Indexed {len(new_files)} documents from {args.input_dir}")
-
-        except Exception as e:
-            logging.error(f"Error during startup indexing: {str(e)}")
-
-    @app.post("/documents/scan", dependencies=[Depends(optional_api_key)])
-    async def scan_for_new_documents():
-        """Manually trigger scanning for new documents"""
-        try:
-            new_files = doc_manager.scan_directory()
-            indexed_count = 0
-
-            for file_path in new_files:
-                try:
-                    with open(file_path, "r", encoding="utf-8") as f:
-                        content = f.read()
-                        rag.insert(content)
-                        doc_manager.mark_as_indexed(file_path)
-                        indexed_count += 1
-                except Exception as e:
-                    logging.error(f"Error indexing file {file_path}: {str(e)}")
-
-            return {
-                "status": "success",
-                "indexed_count": indexed_count,
-                "total_documents": len(doc_manager.indexed_files),
-            }
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post("/documents/upload", dependencies=[Depends(optional_api_key)])
-    async def upload_to_input_dir(file: UploadFile = File(...)):
-        """Upload a file to the input directory"""
-        try:
-            if not doc_manager.is_supported_file(file.filename):
-                raise HTTPException(
-                    status_code=400,
-                    detail=f"Unsupported file type. Supported types: {doc_manager.supported_extensions}",
-                )
-
-            file_path = doc_manager.input_dir / file.filename
-            with open(file_path, "wb") as buffer:
-                shutil.copyfileobj(file.file, buffer)
-
-            # Immediately index the uploaded file
-            with open(file_path, "r", encoding="utf-8") as f:
-                content = f.read()
-                rag.insert(content)
-                doc_manager.mark_as_indexed(file_path)
-
-            return {
-                "status": "success",
-                "message": f"File uploaded and indexed: {file.filename}",
-                "total_documents": len(doc_manager.indexed_files),
-            }
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
-    )
-    async def query_text(request: QueryRequest):
-        try:
-            response = await rag.aquery(
-                request.query,
-                param=QueryParam(
-                    mode=request.mode,
-                    stream=request.stream,
-                    only_need_context=request.only_need_context,
-                ),
-            )
-
-            if request.stream:
-                result = ""
-                async for chunk in response:
-                    result += chunk
-                return QueryResponse(response=result)
-            else:
-                return QueryResponse(response=response)
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post("/query/stream", dependencies=[Depends(optional_api_key)])
-    async def query_text_stream(request: QueryRequest):
-        try:
-            response = rag.query(
-                request.query,
-                param=QueryParam(
-                    mode=request.mode,
-                    stream=True,
-                    only_need_context=request.only_need_context,
-                ),
-            )
-
-            async def stream_generator():
-                async for chunk in response:
-                    yield chunk
-
-            return stream_generator()
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/documents/text",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def insert_text(request: InsertTextRequest):
-        try:
-            rag.insert(request.text)
-            return InsertResponse(
-                status="success",
-                message="Text successfully inserted",
-                document_count=len(rag),
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/documents/file",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
-        try:
-            content = await file.read()
-
-            if file.filename.endswith((".txt", ".md")):
-                text = content.decode("utf-8")
-                rag.insert(text)
-            else:
-                raise HTTPException(
-                    status_code=400,
-                    detail="Unsupported file type. Only .txt and .md files are supported",
-                )
-
-            return InsertResponse(
-                status="success",
-                message=f"File '{file.filename}' successfully inserted",
-                document_count=1,
-            )
-        except UnicodeDecodeError:
-            raise HTTPException(status_code=400, detail="File encoding not supported")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.post(
-        "/documents/batch",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def insert_batch(files: List[UploadFile] = File(...)):
-        try:
-            inserted_count = 0
-            failed_files = []
-
-            for file in files:
-                try:
-                    content = await file.read()
-                    if file.filename.endswith((".txt", ".md")):
-                        text = content.decode("utf-8")
-                        rag.insert(text)
-                        inserted_count += 1
-                    else:
-                        failed_files.append(f"{file.filename} (unsupported type)")
-                except Exception as e:
-                    failed_files.append(f"{file.filename} ({str(e)})")
-
-            status_message = f"Successfully inserted {inserted_count} documents"
-            if failed_files:
-                status_message += f". Failed files: {', '.join(failed_files)}"
-
-            return InsertResponse(
-                status="success" if inserted_count > 0 else "partial_success",
-                message=status_message,
-                document_count=len(files),
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.delete(
-        "/documents",
-        response_model=InsertResponse,
-        dependencies=[Depends(optional_api_key)],
-    )
-    async def clear_documents():
-        try:
-            rag.text_chunks = []
-            rag.entities_vdb = None
-            rag.relationships_vdb = None
-            return InsertResponse(
-                status="success",
-                message="All documents cleared successfully",
-                document_count=0,
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @app.get("/health", dependencies=[Depends(optional_api_key)])
-    async def get_status():
-        """Get current system status"""
-        return {
-            "status": "healthy",
-            "working_directory": str(args.working_dir),
-            "input_directory": str(args.input_dir),
-            "indexed_files": len(doc_manager.indexed_files),
-            "configuration": {
-                "model": args.model,
-                "embedding_model": args.embedding_model,
-                "max_tokens": args.max_tokens,
-                "embedding_dim": embedding_dim,
-            },
-        }
-
-    return app
-
-
-def main():
-    args = parse_args()
-    import uvicorn
-
-    app = create_app(args)
-    uvicorn.run(app, host=args.host, port=args.port)
-
-
-if __name__ == "__main__":
-    main()

lightrag/api/requirements.txt

@@ -7,6 +7,7 @@ nest_asyncio
 numpy
 ollama
 openai
+pipmaster
 python-dotenv
 python-multipart
 tenacity

lightrag/kg/mongo_impl.py

@@ -2,7 +2,7 @@ import os
 from tqdm.asyncio import tqdm as tqdm_async
 from dataclasses import dataclass
 from pymongo import MongoClient
-
+from typing import Union
 from lightrag.utils import logger
 
 from lightrag.base import BaseKVStorage
@@ -41,11 +41,35 @@ class MongoKVStorage(BaseKVStorage):
         return set([s for s in data if s not in existing_ids])
 
     async def upsert(self, data: dict[str, dict]):
-        for k, v in tqdm_async(data.items(), desc="Upserting"):
-            self._data.update_one({"_id": k}, {"$set": v}, upsert=True)
-            data[k]["_id"] = k
+        if self.namespace == "llm_response_cache":
+            for mode, items in data.items():
+                for k, v in tqdm_async(items.items(), desc="Upserting"):
+                    key = f"{mode}_{k}"
+                    result = self._data.update_one(
+                        {"_id": key}, {"$setOnInsert": v}, upsert=True
+                    )
+                    if result.upserted_id:
+                        logger.debug(f"\nInserted new document with key: {key}")
+                    data[mode][k]["_id"] = key
+        else:
+            for k, v in tqdm_async(data.items(), desc="Upserting"):
+                self._data.update_one({"_id": k}, {"$set": v}, upsert=True)
+                data[k]["_id"] = k
         return data
 
+    async def get_by_mode_and_id(self, mode: str, id: str) -> Union[dict, None]:
+        if "llm_response_cache" == self.namespace:
+            res = {}
+            v = self._data.find_one({"_id": mode + "_" + id})
+            if v:
+                res[id] = v
+                logger.debug(f"llm_response_cache find one by:{id}")
+                return res
+            else:
+                return None
+        else:
+            return None
+
     async def drop(self):
         """ """
         pass

lightrag/kg/neo4j_impl.py

@@ -39,6 +39,7 @@ class Neo4JStorage(BaseGraphStorage):
         URI = os.environ["NEO4J_URI"]
         USERNAME = os.environ["NEO4J_USERNAME"]
         PASSWORD = os.environ["NEO4J_PASSWORD"]
+        MAX_CONNECTION_POOL_SIZE = os.environ.get("NEO4J_MAX_CONNECTION_POOL_SIZE", 800)
         DATABASE = os.environ.get(
             "NEO4J_DATABASE"
         )  # If this param is None, the home database will be used. If it is not None, the specified database will be used.
@@ -47,7 +48,11 @@ class Neo4JStorage(BaseGraphStorage):
             URI, auth=(USERNAME, PASSWORD)
         )
         _database_name = "home database" if DATABASE is None else f"database {DATABASE}"
-        with GraphDatabase.driver(URI, auth=(USERNAME, PASSWORD)) as _sync_driver:
+        with GraphDatabase.driver(
+            URI,
+            auth=(USERNAME, PASSWORD),
+            max_connection_pool_size=MAX_CONNECTION_POOL_SIZE,
+        ) as _sync_driver:
             try:
                 with _sync_driver.session(database=DATABASE) as session:
                     try:

lightrag/kg/postgres_impl.py

@@ -130,6 +130,7 @@ class PostgreSQLDB:
         data: Union[list, dict] = None,
         for_age: bool = False,
         graph_name: str = None,
+        upsert: bool = False,
     ):
         try:
             async with self.pool.acquire() as connection:
@@ -140,8 +141,16 @@ class PostgreSQLDB:
                     await connection.execute(sql)
                 else:
                     await connection.execute(sql, *data.values())
+        except (
+            asyncpg.exceptions.UniqueViolationError,
+            asyncpg.exceptions.DuplicateTableError,
+        ) as e:
+            if upsert:
+                print("Key value duplicate, but upsert succeeded.")
+            else:
+                logger.error(f"Upsert error: {e}")
         except Exception as e:
-            logger.error(f"PostgreSQL database error: {e}")
+            logger.error(f"PostgreSQL database error: {e.__class__} - {e}")
             print(sql)
             print(data)
             raise
@@ -568,10 +577,10 @@ class PGGraphStorage(BaseGraphStorage):
 
             if dtype == "vertex":
                 vertex = json.loads(v)
-                field = json.loads(v).get("properties")
+                field = vertex.get("properties")
                 if not field:
                     field = {}
-                field["label"] = PGGraphStorage._decode_graph_label(vertex["label"])
+                field["label"] = PGGraphStorage._decode_graph_label(field["node_id"])
                 d[k] = field
             # convert edge from id-label->id by replacing id with node information
             # we only do this if the vertex was also returned in the query
@@ -666,73 +675,8 @@ class PGGraphStorage(BaseGraphStorage):
         # otherwise return the value stripping out some common special chars
         return field.replace("(", "_").replace(")", "")
 
-    @staticmethod
-    def _wrap_query(query: str, graph_name: str, **params: str) -> str:
-        """
-        Convert a cypher query to an Apache Age compatible
-        sql query by wrapping the cypher query in ag_catalog.cypher,
-        casting results to agtype and building a select statement
-
-        Args:
-            query (str): a valid cypher query
-            graph_name (str): the name of the graph to query
-            params (dict): parameters for the query
-
-        Returns:
-            str: an equivalent pgsql query
-        """
-
-        # pgsql template
-        template = """SELECT {projection} FROM ag_catalog.cypher('{graph_name}', $$
-           {query}
-        $$) AS ({fields})"""
-
-        # if there are any returned fields they must be added to the pgsql query
-        if "return" in query.lower():
-            # parse return statement to identify returned fields
-            fields = (
-                query.lower()
-                .split("return")[-1]
-                .split("distinct")[-1]
-                .split("order by")[0]
-                .split("skip")[0]
-                .split("limit")[0]
-                .split(",")
-            )
-
-            # raise exception if RETURN * is found as we can't resolve the fields
-            if "*" in [x.strip() for x in fields]:
-                raise ValueError(
-                    "AGE graph does not support 'RETURN *'"
-                    + " statements in Cypher queries"
-                )
-
-            # get pgsql formatted field names
-            fields = [
-                PGGraphStorage._get_col_name(field, idx)
-                for idx, field in enumerate(fields)
-            ]
-
-            # build resulting pgsql relation
-            fields_str = ", ".join(
-                [field.split(".")[-1] + " agtype" for field in fields]
-            )
-
-        # if no return statement we still need to return a single field of type agtype
-        else:
-            fields_str = "a agtype"
-
-        select_str = "*"
-
-        return template.format(
-            graph_name=graph_name,
-            query=query.format(**params),
-            fields=fields_str,
-            projection=select_str,
-        )
-
     async def _query(
-        self, query: str, readonly=True, upsert_edge=False, **params: str
+        self, query: str, readonly: bool = True, upsert: bool = False
     ) -> List[Dict[str, Any]]:
         """
         Query the graph by taking a cypher query, converting it to an
@@ -746,7 +690,7 @@ class PGGraphStorage(BaseGraphStorage):
             List[Dict[str, Any]]: a list of dictionaries containing the result set
         """
         # convert cypher query to pgsql/age query
-        wrapped_query = self._wrap_query(query, self.graph_name, **params)
+        wrapped_query = query
 
         # execute the query, rolling back on an error
         try:
@@ -758,22 +702,16 @@ class PGGraphStorage(BaseGraphStorage):
                     graph_name=self.graph_name,
                 )
             else:
-                # for upserting edge, need to run the SQL twice, otherwise cannot update the properties. (First time it will try to create the edge, second time is MERGING)
-                # It is a bug of AGE as of 2025-01-03, hope it can be resolved in the future.
-                if upsert_edge:
-                    data = await self.db.execute(
-                        f"{wrapped_query};{wrapped_query};",
-                        for_age=True,
-                        graph_name=self.graph_name,
-                    )
-                else:
-                    data = await self.db.execute(
-                        wrapped_query, for_age=True, graph_name=self.graph_name
-                    )
+                data = await self.db.execute(
+                    wrapped_query,
+                    for_age=True,
+                    graph_name=self.graph_name,
+                    upsert=upsert,
+                )
         except Exception as e:
             raise PGGraphQueryException(
                 {
-                    "message": f"Error executing graph query: {query.format(**params)}",
+                    "message": f"Error executing graph query: {query}",
                     "wrapped": wrapped_query,
                     "detail": str(e),
                 }
@@ -788,77 +726,85 @@ class PGGraphStorage(BaseGraphStorage):
         return result
 
     async def has_node(self, node_id: str) -> bool:
-        entity_name_label = node_id.strip('"')
+        entity_name_label = PGGraphStorage._encode_graph_label(node_id.strip('"'))
+
+        query = """SELECT * FROM cypher('%s', $$
+                     MATCH (n:Entity {node_id: "%s"})
+                     RETURN count(n) > 0 AS node_exists
+                   $$) AS (node_exists bool)""" % (self.graph_name, entity_name_label)
 
-        query = """MATCH (n:`{label}`) RETURN count(n) > 0 AS node_exists"""
-        params = {"label": PGGraphStorage._encode_graph_label(entity_name_label)}
-        single_result = (await self._query(query, **params))[0]
+        single_result = (await self._query(query))[0]
         logger.debug(
             "{%s}:query:{%s}:result:{%s}",
             inspect.currentframe().f_code.co_name,
-            query.format(**params),
+            query,
             single_result["node_exists"],
         )
 
         return single_result["node_exists"]
 
     async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
-        entity_name_label_source = source_node_id.strip('"')
-        entity_name_label_target = target_node_id.strip('"')
+        src_label = PGGraphStorage._encode_graph_label(source_node_id.strip('"'))
+        tgt_label = PGGraphStorage._encode_graph_label(target_node_id.strip('"'))
+
+        query = """SELECT * FROM cypher('%s', $$
+                     MATCH (a:Entity {node_id: "%s"})-[r]-(b:Entity {node_id: "%s"})
+                     RETURN COUNT(r) > 0 AS edge_exists
+                   $$) AS (edge_exists bool)""" % (
+            self.graph_name,
+            src_label,
+            tgt_label,
+        )
 
-        query = """MATCH (a:`{src_label}`)-[r]-(b:`{tgt_label}`)
-                RETURN COUNT(r) > 0 AS edge_exists"""
-        params = {
-            "src_label": PGGraphStorage._encode_graph_label(entity_name_label_source),
-            "tgt_label": PGGraphStorage._encode_graph_label(entity_name_label_target),
-        }
-        single_result = (await self._query(query, **params))[0]
+        single_result = (await self._query(query))[0]
         logger.debug(
             "{%s}:query:{%s}:result:{%s}",
             inspect.currentframe().f_code.co_name,
-            query.format(**params),
+            query,
             single_result["edge_exists"],
         )
         return single_result["edge_exists"]
 
     async def get_node(self, node_id: str) -> Union[dict, None]:
-        entity_name_label = node_id.strip('"')
-        query = """MATCH (n:`{label}`) RETURN n"""
-        params = {"label": PGGraphStorage._encode_graph_label(entity_name_label)}
-        record = await self._query(query, **params)
+        label = PGGraphStorage._encode_graph_label(node_id.strip('"'))
+        query = """SELECT * FROM cypher('%s', $$
+                     MATCH (n:Entity {node_id: "%s"})
+                     RETURN n
+                   $$) AS (n agtype)""" % (self.graph_name, label)
+        record = await self._query(query)
         if record:
             node = record[0]
             node_dict = node["n"]
             logger.debug(
                 "{%s}: query: {%s}, result: {%s}",
                 inspect.currentframe().f_code.co_name,
-                query.format(**params),
+                query,
                 node_dict,
             )
             return node_dict
         return None
 
     async def node_degree(self, node_id: str) -> int:
-        entity_name_label = node_id.strip('"')
+        label = PGGraphStorage._encode_graph_label(node_id.strip('"'))
 
-        query = """MATCH (n:`{label}`)-[]->(x) RETURN count(x) AS total_edge_count"""
-        params = {"label": PGGraphStorage._encode_graph_label(entity_name_label)}
-        record = (await self._query(query, **params))[0]
+        query = """SELECT * FROM cypher('%s', $$
+                     MATCH (n:Entity {node_id: "%s"})-[]->(x)
+                     RETURN count(x) AS total_edge_count
+                   $$) AS (total_edge_count integer)""" % (self.graph_name, label)
+        record = (await self._query(query))[0]
         if record:
             edge_count = int(record["total_edge_count"])
             logger.debug(
                 "{%s}:query:{%s}:result:{%s}",
                 inspect.currentframe().f_code.co_name,
-                query.format(**params),
+                query,
                 edge_count,
             )
             return edge_count
 
     async def edge_degree(self, src_id: str, tgt_id: str) -> int:
-        entity_name_label_source = src_id.strip('"')
-        entity_name_label_target = tgt_id.strip('"')
-        src_degree = await self.node_degree(entity_name_label_source)
-        trg_degree = await self.node_degree(entity_name_label_target)
+        src_degree = await self.node_degree(src_id)
+        trg_degree = await self.node_degree(tgt_id)
 
         # Convert None to 0 for addition
         src_degree = 0 if src_degree is None else src_degree
@@ -885,23 +831,25 @@ class PGGraphStorage(BaseGraphStorage):
         Returns:
             list: List of all relationships/edges found
         """
-        entity_name_label_source = source_node_id.strip('"')
-        entity_name_label_target = target_node_id.strip('"')
-
-        query = """MATCH (a:`{src_label}`)-[r]->(b:`{tgt_label}`)
-                RETURN properties(r) as edge_properties
-                LIMIT 1"""
-        params = {
-            "src_label": PGGraphStorage._encode_graph_label(entity_name_label_source),
-            "tgt_label": PGGraphStorage._encode_graph_label(entity_name_label_target),
-        }
-        record = await self._query(query, **params)
+        src_label = PGGraphStorage._encode_graph_label(source_node_id.strip('"'))
+        tgt_label = PGGraphStorage._encode_graph_label(target_node_id.strip('"'))
+
+        query = """SELECT * FROM cypher('%s', $$
+                     MATCH (a:Entity {node_id: "%s"})-[r]->(b:Entity {node_id: "%s"})
+                     RETURN properties(r) as edge_properties
+                     LIMIT 1
+                   $$) AS (edge_properties agtype)""" % (
+            self.graph_name,
+            src_label,
+            tgt_label,
+        )
+        record = await self._query(query)
         if record and record[0] and record[0]["edge_properties"]:
             result = record[0]["edge_properties"]
             logger.debug(
                 "{%s}:query:{%s}:result:{%s}",
                 inspect.currentframe().f_code.co_name,
-                query.format(**params),
+                query,
                 result,
             )
             return result
@@ -911,29 +859,41 @@ class PGGraphStorage(BaseGraphStorage):
         Retrieves all edges (relationships) for a particular node identified by its label.
         :return: List of dictionaries containing edge information
         """
-        node_label = source_node_id.strip('"')
+        label = PGGraphStorage._encode_graph_label(source_node_id.strip('"'))
+
+        query = """SELECT * FROM cypher('%s', $$
+                      MATCH (n:Entity {node_id: "%s"})
+                      OPTIONAL MATCH (n)-[r]-(connected)
+                      RETURN n, r, connected
+                    $$) AS (n agtype, r agtype, connected agtype)""" % (
+            self.graph_name,
+            label,
+        )
 
-        query = """MATCH (n:`{label}`)
-                OPTIONAL MATCH (n)-[r]-(connected)
-                RETURN n, r, connected"""
-        params = {"label": PGGraphStorage._encode_graph_label(node_label)}
-        results = await self._query(query, **params)
+        results = await self._query(query)
         edges = []
         for record in results:
             source_node = record["n"] if record["n"] else None
             connected_node = record["connected"] if record["connected"] else None
 
             source_label = (
-                source_node["label"] if source_node and source_node["label"] else None
+                source_node["node_id"]
+                if source_node and source_node["node_id"]
+                else None
             )
             target_label = (
-                connected_node["label"]
-                if connected_node and connected_node["label"]
+                connected_node["node_id"]
+                if connected_node and connected_node["node_id"]
                 else None
             )
 
             if source_label and target_label:
-                edges.append((source_label, target_label))
+                edges.append(
+                    (
+                        PGGraphStorage._decode_graph_label(source_label),
+                        PGGraphStorage._decode_graph_label(target_label),
+                    )
+                )
 
         return edges
 
@@ -950,17 +910,21 @@ class PGGraphStorage(BaseGraphStorage):
             node_id: The unique identifier for the node (used as label)
             node_data: Dictionary of node properties
         """
-        label = node_id.strip('"')
+        label = PGGraphStorage._encode_graph_label(node_id.strip('"'))
         properties = node_data
 
-        query = """MERGE (n:`{label}`)
-                SET n += {properties}"""
-        params = {
-            "label": PGGraphStorage._encode_graph_label(label),
-            "properties": PGGraphStorage._format_properties(properties),
-        }
+        query = """SELECT * FROM cypher('%s', $$
+                     MERGE (n:Entity {node_id: "%s"})
+                     SET n += %s
+                     RETURN n
+                   $$) AS (n agtype)""" % (
+            self.graph_name,
+            label,
+            PGGraphStorage._format_properties(properties),
+        )
+
         try:
-            await self._query(query, readonly=False, **params)
+            await self._query(query, readonly=False, upsert=True)
             logger.debug(
                 "Upserted node with label '{%s}' and properties: {%s}",
                 label,
@@ -986,28 +950,30 @@ class PGGraphStorage(BaseGraphStorage):
             target_node_id (str): Label of the target node (used as identifier)
             edge_data (dict): Dictionary of properties to set on the edge
         """
-        source_node_label = source_node_id.strip('"')
-        target_node_label = target_node_id.strip('"')
+        src_label = PGGraphStorage._encode_graph_label(source_node_id.strip('"'))
+        tgt_label = PGGraphStorage._encode_graph_label(target_node_id.strip('"'))
         edge_properties = edge_data
 
-        query = """MATCH (source:`{src_label}`)
-                WITH source
-                MATCH (target:`{tgt_label}`)
-                MERGE (source)-[r:DIRECTED]->(target)
-                SET r += {properties}
-                RETURN r"""
-        params = {
-            "src_label": PGGraphStorage._encode_graph_label(source_node_label),
-            "tgt_label": PGGraphStorage._encode_graph_label(target_node_label),
-            "properties": PGGraphStorage._format_properties(edge_properties),
-        }
+        query = """SELECT * FROM cypher('%s', $$
+                     MATCH (source:Entity {node_id: "%s"})
+                     WITH source
+                     MATCH (target:Entity {node_id: "%s"})
+                     MERGE (source)-[r:DIRECTED]->(target)
+                     SET r += %s
+                     RETURN r
+                   $$) AS (r agtype)""" % (
+            self.graph_name,
+            src_label,
+            tgt_label,
+            PGGraphStorage._format_properties(edge_properties),
+        )
         # logger.info(f"-- inserting edge after formatted: {params}")
         try:
-            await self._query(query, readonly=False, upsert_edge=True, **params)
+            await self._query(query, readonly=False, upsert=True)
             logger.debug(
                 "Upserted edge from '{%s}' to '{%s}' with properties: {%s}",
-                source_node_label,
-                target_node_label,
+                src_label,
+                tgt_label,
                 edge_properties,
             )
         except Exception as e:

lightrag/kg/postgres_impl_test.py

@@ -61,7 +61,7 @@ db = PostgreSQLDB(
         "port": 15432,
         "user": "rag",
         "password": "rag",
-        "database": "rag",
+        "database": "r1",
     }
 )
 
@@ -74,8 +74,12 @@ async def query_with_age():
         embedding_func=None,
     )
     graph.db = db
-    res = await graph.get_node('"CHRISTMAS-TIME"')
+    res = await graph.get_node('"A CHRISTMAS CAROL"')
     print("Node is: ", res)
+    res = await graph.get_edge('"A CHRISTMAS CAROL"', "PROJECT GUTENBERG")
+    print("Edge is: ", res)
+    res = await graph.get_node_edges('"SCROOGE"')
+    print("Node Edges are: ", res)
 
 
 async def create_edge_with_age():

lightrag/lightrag.py

@@ -45,6 +45,7 @@ from .storage import (
 
 from .prompt import GRAPH_FIELD_SEP
 
+
 # future KG integrations
 
 # from .kg.ArangoDB_impl import (
@@ -168,7 +169,7 @@ class LightRAG:
 
     # LLM
     llm_model_func: callable = gpt_4o_mini_complete  # hf_model_complete#
-    llm_model_name: str = "meta-llama/Llama-3.2-1B-Instruct"  #'meta-llama/Llama-3.2-1B'#'google/gemma-2-2b-it'
+    llm_model_name: str = "meta-llama/Llama-3.2-1B-Instruct"  # 'meta-llama/Llama-3.2-1B'#'google/gemma-2-2b-it'
     llm_model_max_token_size: int = 32768
     llm_model_max_async: int = 16
     llm_model_kwargs: dict = field(default_factory=dict)
@@ -187,6 +188,10 @@ class LightRAG:
     # Add new field for document status storage type
     doc_status_storage: str = field(default="JsonDocStatusStorage")
 
+    # Custom Chunking Function
+    chunking_func: callable = chunking_by_token_size
+    chunking_func_kwargs: dict = field(default_factory=dict)
+
     def __post_init__(self):
         log_file = os.path.join("lightrag.log")
         set_logger(log_file)
@@ -315,15 +320,25 @@ class LightRAG:
             "JsonDocStatusStorage": JsonDocStatusStorage,
         }
 
-    def insert(self, string_or_strings):
+    def insert(
+        self, string_or_strings, split_by_character=None, split_by_character_only=False
+    ):
         loop = always_get_an_event_loop()
-        return loop.run_until_complete(self.ainsert(string_or_strings))
+        return loop.run_until_complete(
+            self.ainsert(string_or_strings, split_by_character, split_by_character_only)
+        )
 
-    async def ainsert(self, string_or_strings):
+    async def ainsert(
+        self, string_or_strings, split_by_character=None, split_by_character_only=False
+    ):
         """Insert documents with checkpoint support
 
         Args:
             string_or_strings: Single document string or list of document strings
+            split_by_character: if split_by_character is not None, split the string by character, if chunk longer than
+            chunk_size, split the sub chunk by token size.
+            split_by_character_only: if split_by_character_only is True, split the string by character only, when
+            split_by_character is None, this parameter is ignored.
         """
         if isinstance(string_or_strings, str):
             string_or_strings = [string_or_strings]
@@ -360,7 +375,7 @@ class LightRAG:
             batch_docs = dict(list(new_docs.items())[i : i + batch_size])
 
             for doc_id, doc in tqdm_async(
-                batch_docs.items(), desc=f"Processing batch {i//batch_size + 1}"
+                batch_docs.items(), desc=f"Processing batch {i // batch_size + 1}"
             ):
                 try:
                     # Update status to processing
@@ -379,11 +394,14 @@ class LightRAG:
                             **dp,
                             "full_doc_id": doc_id,
                         }
-                        for dp in chunking_by_token_size(
+                        for dp in self.chunking_func(
                             doc["content"],
+                            split_by_character=split_by_character,
+                            split_by_character_only=split_by_character_only,
                             overlap_token_size=self.chunk_overlap_token_size,
                             max_token_size=self.chunk_token_size,
                             tiktoken_model=self.tiktoken_model_name,
+                            **self.chunking_func_kwargs,
                         )
                     }
 
@@ -455,6 +473,73 @@ class LightRAG:
                     # Ensure all indexes are updated after each document
                     await self._insert_done()
 
+    def insert_custom_chunks(self, full_text: str, text_chunks: list[str]):
+        loop = always_get_an_event_loop()
+        return loop.run_until_complete(
+            self.ainsert_custom_chunks(full_text, text_chunks)
+        )
+
+    async def ainsert_custom_chunks(self, full_text: str, text_chunks: list[str]):
+        update_storage = False
+        try:
+            doc_key = compute_mdhash_id(full_text.strip(), prefix="doc-")
+            new_docs = {doc_key: {"content": full_text.strip()}}
+
+            _add_doc_keys = await self.full_docs.filter_keys([doc_key])
+            new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}
+            if not len(new_docs):
+                logger.warning("This document is already in the storage.")
+                return
+
+            update_storage = True
+            logger.info(f"[New Docs] inserting {len(new_docs)} docs")
+
+            inserting_chunks = {}
+            for chunk_text in text_chunks:
+                chunk_text_stripped = chunk_text.strip()
+                chunk_key = compute_mdhash_id(chunk_text_stripped, prefix="chunk-")
+
+                inserting_chunks[chunk_key] = {
+                    "content": chunk_text_stripped,
+                    "full_doc_id": doc_key,
+                }
+
+            _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("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)
+
+            logger.info("[Entity Extraction]...")
+            maybe_new_kg = await extract_entities(
+                inserting_chunks,
+                knowledge_graph_inst=self.chunk_entity_relation_graph,
+                entity_vdb=self.entities_vdb,
+                relationships_vdb=self.relationships_vdb,
+                global_config=asdict(self),
+            )
+
+            if maybe_new_kg is None:
+                logger.warning("No new entities and relationships found")
+                return
+            else:
+                self.chunk_entity_relation_graph = maybe_new_kg
+
+            await self.full_docs.upsert(new_docs)
+            await self.text_chunks.upsert(inserting_chunks)
+
+        finally:
+            if update_storage:
+                await self._insert_done()
+
     async def _insert_done(self):
         tasks = []
         for storage_inst in [

lightrag/llm.py

@@ -406,8 +406,9 @@ async def lollms_model_if_cache(
     full_prompt += prompt
 
     request_data["prompt"] = full_prompt
+    timeout = aiohttp.ClientTimeout(total=kwargs.get("timeout", None))
 
-    async with aiohttp.ClientSession() as session:
+    async with aiohttp.ClientSession(timeout=timeout) as session:
         if stream:
 
             async def inner():

lightrag/operate.py

@@ -4,7 +4,6 @@ import re
 from tqdm.asyncio import tqdm as tqdm_async
 from typing import Union
 from collections import Counter, defaultdict
-import warnings
 from .utils import (
     logger,
     clean_str,
@@ -34,23 +33,61 @@ import time
 
 
 def chunking_by_token_size(
-    content: str, overlap_token_size=128, max_token_size=1024, tiktoken_model="gpt-4o"
+    content: str,
+    split_by_character=None,
+    split_by_character_only=False,
+    overlap_token_size=128,
+    max_token_size=1024,
+    tiktoken_model="gpt-4o",
+    **kwargs,
 ):
     tokens = encode_string_by_tiktoken(content, model_name=tiktoken_model)
     results = []
-    for index, start in enumerate(
-        range(0, len(tokens), max_token_size - overlap_token_size)
-    ):
-        chunk_content = decode_tokens_by_tiktoken(
-            tokens[start : start + max_token_size], model_name=tiktoken_model
-        )
-        results.append(
-            {
-                "tokens": min(max_token_size, len(tokens) - start),
-                "content": chunk_content.strip(),
-                "chunk_order_index": index,
-            }
-        )
+    if split_by_character:
+        raw_chunks = content.split(split_by_character)
+        new_chunks = []
+        if split_by_character_only:
+            for chunk in raw_chunks:
+                _tokens = encode_string_by_tiktoken(chunk, model_name=tiktoken_model)
+                new_chunks.append((len(_tokens), chunk))
+        else:
+            for chunk in raw_chunks:
+                _tokens = encode_string_by_tiktoken(chunk, model_name=tiktoken_model)
+                if len(_tokens) > max_token_size:
+                    for start in range(
+                        0, len(_tokens), max_token_size - overlap_token_size
+                    ):
+                        chunk_content = decode_tokens_by_tiktoken(
+                            _tokens[start : start + max_token_size],
+                            model_name=tiktoken_model,
+                        )
+                        new_chunks.append(
+                            (min(max_token_size, len(_tokens) - start), chunk_content)
+                        )
+                else:
+                    new_chunks.append((len(_tokens), chunk))
+        for index, (_len, chunk) in enumerate(new_chunks):
+            results.append(
+                {
+                    "tokens": _len,
+                    "content": chunk.strip(),
+                    "chunk_order_index": index,
+                }
+            )
+    else:
+        for index, start in enumerate(
+            range(0, len(tokens), max_token_size - overlap_token_size)
+        ):
+            chunk_content = decode_tokens_by_tiktoken(
+                tokens[start : start + max_token_size], model_name=tiktoken_model
+            )
+            results.append(
+                {
+                    "tokens": min(max_token_size, len(tokens) - start),
+                    "content": chunk_content.strip(),
+                    "chunk_order_index": index,
+                }
+            )
     return results
 
 
@@ -582,15 +619,22 @@ async def kg_query(
         logger.warning("low_level_keywords and high_level_keywords is empty")
         return PROMPTS["fail_response"]
     if ll_keywords == [] and query_param.mode in ["local", "hybrid"]:
-        logger.warning("low_level_keywords is empty")
-        return PROMPTS["fail_response"]
-    else:
-        ll_keywords = ", ".join(ll_keywords)
+        logger.warning(
+            "low_level_keywords is empty, switching from %s mode to global mode",
+            query_param.mode,
+        )
+        query_param.mode = "global"
     if hl_keywords == [] and query_param.mode in ["global", "hybrid"]:
-        logger.warning("high_level_keywords is empty")
-        return PROMPTS["fail_response"]
-    else:
-        hl_keywords = ", ".join(hl_keywords)
+        logger.warning(
+            "high_level_keywords is empty, switching from %s mode to local mode",
+            query_param.mode,
+        )
+        query_param.mode = "local"
+
+    ll_keywords = ", ".join(ll_keywords) if ll_keywords else ""
+    hl_keywords = ", ".join(hl_keywords) if hl_keywords else ""
+
+    logger.info("Using %s mode for query processing", query_param.mode)
 
     # Build context
     keywords = [ll_keywords, hl_keywords]
@@ -656,77 +700,51 @@ async def _build_query_context(
     # ll_entities_context, ll_relations_context, ll_text_units_context = "", "", ""
     # hl_entities_context, hl_relations_context, hl_text_units_context = "", "", ""
 
-    ll_kewwords, hl_keywrds = query[0], query[1]
-    if query_param.mode in ["local", "hybrid"]:
-        if ll_kewwords == "":
-            ll_entities_context, ll_relations_context, ll_text_units_context = (
-                "",
-                "",
-                "",
-            )
-            warnings.warn(
-                "Low Level context is None. Return empty Low entity/relationship/source"
-            )
-            query_param.mode = "global"
-        else:
-            (
-                ll_entities_context,
-                ll_relations_context,
-                ll_text_units_context,
-            ) = await _get_node_data(
-                ll_kewwords,
-                knowledge_graph_inst,
-                entities_vdb,
-                text_chunks_db,
-                query_param,
-            )
-    if query_param.mode in ["global", "hybrid"]:
-        if hl_keywrds == "":
-            hl_entities_context, hl_relations_context, hl_text_units_context = (
-                "",
-                "",
-                "",
-            )
-            warnings.warn(
-                "High Level context is None. Return empty High entity/relationship/source"
-            )
-            query_param.mode = "local"
-        else:
-            (
-                hl_entities_context,
-                hl_relations_context,
-                hl_text_units_context,
-            ) = await _get_edge_data(
-                hl_keywrds,
-                knowledge_graph_inst,
-                relationships_vdb,
-                text_chunks_db,
-                query_param,
-            )
-            if (
-                hl_entities_context == ""
-                and hl_relations_context == ""
-                and hl_text_units_context == ""
-            ):
-                logger.warn("No high level context found. Switching to local mode.")
-                query_param.mode = "local"
-    if query_param.mode == "hybrid":
-        entities_context, relations_context, text_units_context = combine_contexts(
-            [hl_entities_context, ll_entities_context],
-            [hl_relations_context, ll_relations_context],
-            [hl_text_units_context, ll_text_units_context],
+    ll_keywords, hl_keywords = query[0], query[1]
+
+    if query_param.mode == "local":
+        entities_context, relations_context, text_units_context = await _get_node_data(
+            ll_keywords,
+            knowledge_graph_inst,
+            entities_vdb,
+            text_chunks_db,
+            query_param,
+        )
+    elif query_param.mode == "global":
+        entities_context, relations_context, text_units_context = await _get_edge_data(
+            hl_keywords,
+            knowledge_graph_inst,
+            relationships_vdb,
+            text_chunks_db,
+            query_param,
         )
-    elif query_param.mode == "local":
-        entities_context, relations_context, text_units_context = (
+    else:  # hybrid mode
+        (
             ll_entities_context,
             ll_relations_context,
             ll_text_units_context,
+        ) = await _get_node_data(
+            ll_keywords,
+            knowledge_graph_inst,
+            entities_vdb,
+            text_chunks_db,
+            query_param,
         )
-    elif query_param.mode == "global":
-        entities_context, relations_context, text_units_context = (
+        (
             hl_entities_context,
             hl_relations_context,
             hl_text_units_context,
+        ) = await _get_edge_data(
+            hl_keywords,
+            knowledge_graph_inst,
+            relationships_vdb,
+            text_chunks_db,
+            query_param,
+        )
+        entities_context, relations_context, text_units_context = combine_contexts(
+            [hl_entities_context, ll_entities_context],
+            [hl_relations_context, ll_relations_context],
+            [hl_text_units_context, ll_text_units_context],
         )
     return f"""
 -----Entities-----

requirements.txt

@@ -1,38 +1,38 @@
 accelerate
-aioboto3~=13.3.0
-aiofiles~=24.1.0
-aiohttp~=3.11.11
-asyncpg~=0.30.0
+aioboto3
+aiofiles
+aiohttp
+asyncpg
 
 # database packages
 graspologic
 gremlinpython
 hnswlib
 nano-vectordb
-neo4j~=5.27.0
-networkx~=3.2.1
+neo4j
+networkx
 
-numpy~=2.2.0
-ollama~=0.4.4
-openai~=1.58.1
+numpy
+ollama
+openai
 oracledb
-psycopg-pool~=3.2.4
-psycopg[binary,pool]~=3.2.3
-pydantic~=2.10.4
+psycopg-pool
+psycopg[binary,pool]
+pydantic
 pymilvus
 pymongo
 pymysql
-python-dotenv~=1.0.1
-pyvis~=0.3.2
-setuptools~=70.0.0
+python-dotenv
+pyvis
+setuptools
 # lmdeploy[all]
-sqlalchemy~=2.0.36
-tenacity~=9.0.0
+sqlalchemy
+tenacity
 
 
 # LLM packages
-tiktoken~=0.8.0
-torch~=2.5.1+cu121
-tqdm~=4.67.1
-transformers~=4.47.1
+tiktoken
+torch
+tqdm
+transformers
 xxhash

setup.py

@@ -100,10 +100,7 @@ setuptools.setup(
     },
     entry_points={
         "console_scripts": [
-            "lollms-lightrag-server=lightrag.api.lollms_lightrag_server:main [api]",
-            "ollama-lightrag-server=lightrag.api.ollama_lightrag_server:main [api]",
-            "openai-lightrag-server=lightrag.api.openai_lightrag_server:main [api]",
-            "azure-openai-lightrag-server=lightrag.api.azure_openai_lightrag_server:main [api]",
+            "lightrag-server=lightrag.api.lightrag_server:main [api]",
         ],
     },
 )