diff --git a/README.md b/README.md index f66fb3ce..71248056 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@

- +

@@ -26,7 +26,8 @@ This repository hosts the code of LightRAG. The structure of this code is based ## ๐ŸŽ‰ 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 - -| 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 | -| --lollms-host | http://localhost:9600 | LoLLMS 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 | - -#### 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 - -| 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 | - -#### OpenAI AZURE Server Options +#### LightRag Server Options | 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 | +| --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 | +| --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-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 | +| --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) | + For protecting the server using an authentication key, you can also use an environment variable named `LIGHTRAG_API_KEY`. ### Example Usage -#### LoLLMs RAG Server +#### Running a Lightrag server with ollama default local server as llm and embedding backends + +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. ```bash -# Custom configuration with specific model and working directory -lollms-lightrag-server --model mistral-nemo --port 8080 --working-dir ./custom_rag +# 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 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-model adrienbrault/nous-hermes2theta-llama3-8b:f16 --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 ollama for embedding +lightrag-server --llm-binding lollms ``` -#### Ollama RAG Server +#### Running a Lightrag server with lollms default local 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, 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 Ollama instance) -ollama-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 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 ``` -#### OpenAI RAG Server + +#### Running a Lightrag server with 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 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 lollms for llm and openai for embedding +lightrag-server --llm-binding lollms --embedding-binding openai --embedding-model text-embedding-3-small ``` +#### Running a Lightrag server with azure openai server as llm and embedding backends + +```bash +# 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 diff --git a/contributor-readme.MD b/contributor-README.md similarity index 100% rename from contributor-readme.MD rename to contributor-README.md diff --git a/examples/test_split_by_character.ipynb b/examples/test_split_by_character.ipynb new file mode 100644 index 00000000..e8e08b92 --- /dev/null +++ b/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ๆ ‡็ญพไธญ๏ผŒ้’ˆๅฏนๆฏไธช้—ฎ้ข˜่ฏฆ็ป†ๅˆ†ๆžไฝ ็š„ๆ€่€ƒ่ฟ‡็จ‹ใ€‚็„ถๅŽๅœจ<ๅ›ž็ญ”>ๆ ‡็ญพไธญ็ป™ๅ‡บๆ‰€ๆœ‰้—ฎ้ข˜็š„ๆœ€็ปˆ็ญ”ๆกˆใ€‚\"\"\"" + ] + }, + { + "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\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 +} diff --git a/lightrag/__init__.py b/lightrag/__init__.py index b8037813..7a26a282 100644 --- a/lightrag/__init__.py +++ b/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" diff --git a/lightrag/api/azure_openai_lightrag_server.py b/lightrag/api/azure_openai_lightrag_server.py deleted file mode 100644 index abe3f738..00000000 --- a/lightrag/api/azure_openai_lightrag_server.py +++ /dev/null @@ -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() diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py new file mode 100644 index 00000000..0d154b38 --- /dev/null +++ b/lightrag/api/lightrag_server.py @@ -0,0 +1,842 @@ +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 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, Union +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 +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(): + parser = argparse.ArgumentParser( + 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)" + ) + 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)", + ) + + # 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( + "--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( + "--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)" + ) + 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, + ) + + # 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", ".pdf", ".docx", ".pptx"), + ): + 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): + # 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) + ) + + # 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.2", + 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=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.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.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""" + try: + new_files = doc_manager.scan_directory() + for file_path in new_files: + try: + await index_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: + await index_file(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 + await index_file(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=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)): + """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 = "" + # 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}") + + return InsertResponse( + status="success", + message=f"File '{file.filename}' successfully inserted", + document_count=1, + ) + else: + raise HTTPException( + status_code=400, + detail="No content could be extracted from the file", + ) + + 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( + "/documents/batch", + response_model=InsertResponse, + 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 = "" + 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} (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)})") + 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=status, + message=status_message, + 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( + "/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": { + # 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, + }, + } + + return app + + +def main(): + args = parse_args() + import uvicorn + + app = create_app(args) + 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__": + main() diff --git a/lightrag/api/lollms_lightrag_server.py b/lightrag/api/lollms_lightrag_server.py deleted file mode 100644 index 8a2804a0..00000000 --- a/lightrag/api/lollms_lightrag_server.py +++ /dev/null @@ -1,492 +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 lollms_model_complete, lollms_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( - "--lollms-host", - default="http://localhost:9600", - help="lollms host URL (default: http://localhost:9600)", - ) - - # 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=lollms_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.lollms_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: lollms_embed( - texts, embed_model=args.embedding_model, host=args.lollms_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: - 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") - 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, - "lollms_host": args.lollms_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() diff --git a/lightrag/api/ollama_lightrag_server.py b/lightrag/api/ollama_lightrag_server.py deleted file mode 100644 index b3140aba..00000000 --- a/lightrag/api/ollama_lightrag_server.py +++ /dev/null @@ -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() diff --git a/lightrag/api/openai_lightrag_server.py b/lightrag/api/openai_lightrag_server.py deleted file mode 100644 index 349c09da..00000000 --- a/lightrag/api/openai_lightrag_server.py +++ /dev/null @@ -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() diff --git a/lightrag/api/requirements.txt b/lightrag/api/requirements.txt index 221d7f40..9154809c 100644 --- a/lightrag/api/requirements.txt +++ b/lightrag/api/requirements.txt @@ -7,6 +7,7 @@ nest_asyncio numpy ollama openai +pipmaster python-dotenv python-multipart tenacity diff --git a/lightrag/kg/mongo_impl.py b/lightrag/kg/mongo_impl.py index 61222357..fbbae8c2 100644 --- a/lightrag/kg/mongo_impl.py +++ b/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 diff --git a/lightrag/kg/neo4j_impl.py b/lightrag/kg/neo4j_impl.py index 884fcb40..8c2afb5d 100644 --- a/lightrag/kg/neo4j_impl.py +++ b/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: diff --git a/lightrag/kg/postgres_impl.py b/lightrag/kg/postgres_impl.py index 033d63d6..b93a345b 100644 --- a/lightrag/kg/postgres_impl.py +++ b/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 = """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] + 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) + + 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 = """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] + 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, + ) + + 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('"') + src_label = PGGraphStorage._encode_graph_label(source_node_id.strip('"')) + tgt_label = PGGraphStorage._encode_graph_label(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) + 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 = """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) + 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, + ) + + 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: diff --git a/lightrag/kg/postgres_impl_test.py b/lightrag/kg/postgres_impl_test.py index dc046311..274f03de 100644 --- a/lightrag/kg/postgres_impl_test.py +++ b/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(): diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index b6d5238e..6e2f1c0e 100644 --- a/lightrag/lightrag.py +++ b/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 [ diff --git a/lightrag/llm.py b/lightrag/llm.py index 0c17019a..c49ed138 100644 --- a/lightrag/llm.py +++ b/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(): diff --git a/lightrag/operate.py b/lightrag/operate.py index 45ba9656..97ac245c 100644 --- a/lightrag/operate.py +++ b/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,78 +700,52 @@ 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": + 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, + ) + 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, + ) + ( + 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], ) - elif query_param.mode == "local": - entities_context, relations_context, text_units_context = ( - ll_entities_context, - ll_relations_context, - ll_text_units_context, - ) - elif query_param.mode == "global": - entities_context, relations_context, text_units_context = ( - hl_entities_context, - hl_relations_context, - hl_text_units_context, - ) return f""" -----Entities----- ```csv diff --git a/requirements.txt b/requirements.txt index 79249e7e..48c25ff8 100644 --- a/requirements.txt +++ b/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 diff --git a/setup.py b/setup.py index 368610f6..38eff646 100644 --- a/setup.py +++ b/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]", ], }, )