ParisNeo commited on
Commit
51da218
·
1 Parent(s): cbf4227

Added openai api

Browse files
api/{README.md → README_OLLAMA.md} RENAMED
File without changes
api/README_OPENAI.md ADDED
@@ -0,0 +1,172 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ # LightRAG API Server
3
+
4
+ A powerful FastAPI-based server for managing and querying documents using LightRAG (Light Retrieval-Augmented Generation). This server provides a REST API interface for document management and intelligent querying using OpenAI's language models.
5
+
6
+ ## Features
7
+
8
+ - 🔍 Multiple search modes (naive, local, global, hybrid)
9
+ - 📡 Streaming and non-streaming responses
10
+ - 📝 Document management (insert, batch upload, clear)
11
+ - ⚙️ Highly configurable model parameters
12
+ - 📚 Support for text and file uploads
13
+ - 🔧 RESTful API with automatic documentation
14
+ - 🚀 Built with FastAPI for high performance
15
+
16
+ ## Prerequisites
17
+
18
+ - Python 3.8+
19
+ - OpenAI API key
20
+ - Required Python packages:
21
+ - fastapi
22
+ - uvicorn
23
+ - lightrag
24
+ - pydantic
25
+ - openai
26
+ - nest-asyncio
27
+
28
+ ## Installation
29
+ If you are using Windows, you will need to download and install visual c++ build tools from [https://visualstudio.microsoft.com/visual-cpp-build-tools/](https://visualstudio.microsoft.com/visual-cpp-build-tools/)
30
+ Make sure you install the VS 2022 C++ x64/x86 Build tools from individual components tab.
31
+
32
+ 1. Clone the repository:
33
+ ```bash
34
+ git clone https://github.com/ParisNeo/LightRAG.git
35
+ cd api
36
+ ```
37
+
38
+ 2. Install dependencies:
39
+ ```bash
40
+ pip install -r requirements.txt
41
+ ```
42
+
43
+ 3. Set up your OpenAI API key:
44
+ ```bash
45
+ export OPENAI_API_KEY='your-api-key-here'
46
+ ```
47
+
48
+ ## Configuration
49
+
50
+ The server can be configured using command-line arguments:
51
+
52
+ ```bash
53
+ python openai_lightrag_server.py --help
54
+ ```
55
+
56
+ Available options:
57
+
58
+ | Parameter | Default | Description |
59
+ |-----------|---------|-------------|
60
+ | --host | 0.0.0.0 | Server host |
61
+ | --port | 9621 | Server port |
62
+ | --model | gpt-4 | OpenAI model name |
63
+ | --embedding-model | text-embedding-3-large | OpenAI embedding model |
64
+ | --working-dir | ./rag_storage | Working directory for RAG |
65
+ | --max-tokens | 32768 | Maximum token size |
66
+ | --max-embed-tokens | 8192 | Maximum embedding token size |
67
+ | --input-dir | ./inputs | Input directory for documents |
68
+ | --log-level | INFO | Logging level |
69
+
70
+ ## Quick Start
71
+
72
+ 1. Basic usage with default settings:
73
+ ```bash
74
+ python openai_lightrag_server.py
75
+ ```
76
+
77
+ 2. Custom configuration:
78
+ ```bash
79
+ python openai_lightrag_server.py --model gpt-4 --port 8080 --working-dir ./custom_rag
80
+ ```
81
+
82
+ ## API Endpoints
83
+
84
+ ### Query Endpoints
85
+
86
+ #### POST /query
87
+ Query the RAG system with options for different search modes.
88
+
89
+ ```bash
90
+ curl -X POST "http://localhost:9621/query" \
91
+ -H "Content-Type: application/json" \
92
+ -d '{"query": "Your question here", "mode": "hybrid"}'
93
+ ```
94
+
95
+ #### POST /query/stream
96
+ Stream responses from the RAG system.
97
+
98
+ ```bash
99
+ curl -X POST "http://localhost:9621/query/stream" \
100
+ -H "Content-Type: application/json" \
101
+ -d '{"query": "Your question here", "mode": "hybrid"}'
102
+ ```
103
+
104
+ ### Document Management Endpoints
105
+
106
+ #### POST /documents/text
107
+ Insert text directly into the RAG system.
108
+
109
+ ```bash
110
+ curl -X POST "http://localhost:9621/documents/text" \
111
+ -H "Content-Type: application/json" \
112
+ -d '{"text": "Your text content here", "description": "Optional description"}'
113
+ ```
114
+
115
+ #### POST /documents/file
116
+ Upload a single file to the RAG system.
117
+
118
+ ```bash
119
+ curl -X POST "http://localhost:9621/documents/file" \
120
+ -F "file=@/path/to/your/document.txt" \
121
+ -F "description=Optional description"
122
+ ```
123
+
124
+ #### POST /documents/batch
125
+ Upload multiple files at once.
126
+
127
+ ```bash
128
+ curl -X POST "http://localhost:9621/documents/batch" \
129
+ -F "files=@/path/to/doc1.txt" \
130
+ -F "files=@/path/to/doc2.txt"
131
+ ```
132
+
133
+ #### DELETE /documents
134
+ Clear all documents from the RAG system.
135
+
136
+ ```bash
137
+ curl -X DELETE "http://localhost:9621/documents"
138
+ ```
139
+
140
+ ### Utility Endpoints
141
+
142
+ #### GET /health
143
+ Check server health and configuration.
144
+
145
+ ```bash
146
+ curl "http://localhost:9621/health"
147
+ ```
148
+
149
+ ## Development
150
+
151
+ ### Running in Development Mode
152
+
153
+ ```bash
154
+ uvicorn openai_lightrag_server:app --reload --port 9621
155
+ ```
156
+
157
+ ### API Documentation
158
+
159
+ When the server is running, visit:
160
+ - Swagger UI: http://localhost:9621/docs
161
+ - ReDoc: http://localhost:9621/redoc
162
+
163
+ ## License
164
+
165
+ This project is licensed under the MIT License - see the LICENSE file for details.
166
+
167
+ ## Acknowledgments
168
+
169
+ - Built with [FastAPI](https://fastapi.tiangolo.com/)
170
+ - Uses [LightRAG](https://github.com/HKUDS/LightRAG) for document processing
171
+ - Powered by [OpenAI](https://openai.com/) for language model inference
172
+
api/openai_lightrag_server.py ADDED
@@ -0,0 +1,369 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, HTTPException, File, UploadFile, Form
2
+ from fastapi.responses import JSONResponse
3
+ from pydantic import BaseModel
4
+ import asyncio
5
+ import os
6
+ import logging
7
+ import argparse
8
+ from lightrag import LightRAG, QueryParam
9
+ from lightrag.llm import openai_complete_if_cache, openai_embedding
10
+ from lightrag.utils import EmbeddingFunc
11
+ from typing import Optional, List
12
+ from enum import Enum
13
+ from pathlib import Path
14
+ import shutil
15
+ import aiofiles
16
+ from ascii_colors import ASCIIColors, trace_exception
17
+ import numpy as np
18
+ import nest_asyncio
19
+
20
+ # Apply nest_asyncio to solve event loop issues
21
+ nest_asyncio.apply()
22
+
23
+ def parse_args():
24
+ parser = argparse.ArgumentParser(
25
+ description="LightRAG FastAPI Server with OpenAI integration"
26
+ )
27
+
28
+ # Server configuration
29
+ parser.add_argument('--host', default='0.0.0.0', help='Server host (default: 0.0.0.0)')
30
+ parser.add_argument('--port', type=int, default=9621, help='Server port (default: 9621)')
31
+
32
+ # Directory configuration
33
+ parser.add_argument('--working-dir', default='./rag_storage',
34
+ help='Working directory for RAG storage (default: ./rag_storage)')
35
+ parser.add_argument('--input-dir', default='./inputs',
36
+ help='Directory containing input documents (default: ./inputs)')
37
+
38
+ # Model configuration
39
+ parser.add_argument('--model', default='gpt-4', help='OpenAI model name (default: gpt-4)')
40
+ parser.add_argument('--embedding-model', default='text-embedding-3-large',
41
+ help='OpenAI embedding model (default: text-embedding-3-large)')
42
+
43
+ # RAG configuration
44
+ parser.add_argument('--max-tokens', type=int, default=32768, help='Maximum token size (default: 32768)')
45
+ parser.add_argument('--max-embed-tokens', type=int, default=8192,
46
+ help='Maximum embedding token size (default: 8192)')
47
+
48
+ # Logging configuration
49
+ parser.add_argument('--log-level', default='INFO',
50
+ choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
51
+ help='Logging level (default: INFO)')
52
+
53
+ return parser.parse_args()
54
+
55
+ class DocumentManager:
56
+ """Handles document operations and tracking"""
57
+
58
+ def __init__(self, input_dir: str, supported_extensions: tuple = ('.txt', '.md')):
59
+ self.input_dir = Path(input_dir)
60
+ self.supported_extensions = supported_extensions
61
+ self.indexed_files = set()
62
+
63
+ # Create input directory if it doesn't exist
64
+ self.input_dir.mkdir(parents=True, exist_ok=True)
65
+
66
+ def scan_directory(self) -> List[Path]:
67
+ """Scan input directory for new files"""
68
+ new_files = []
69
+ for ext in self.supported_extensions:
70
+ for file_path in self.input_dir.rglob(f'*{ext}'):
71
+ if file_path not in self.indexed_files:
72
+ new_files.append(file_path)
73
+ return new_files
74
+
75
+ def mark_as_indexed(self, file_path: Path):
76
+ """Mark a file as indexed"""
77
+ self.indexed_files.add(file_path)
78
+
79
+ def is_supported_file(self, filename: str) -> bool:
80
+ """Check if file type is supported"""
81
+ return any(filename.lower().endswith(ext) for ext in self.supported_extensions)
82
+
83
+ # Pydantic models
84
+ class SearchMode(str, Enum):
85
+ naive = "naive"
86
+ local = "local"
87
+ global_ = "global"
88
+ hybrid = "hybrid"
89
+
90
+ class QueryRequest(BaseModel):
91
+ query: str
92
+ mode: SearchMode = SearchMode.hybrid
93
+ stream: bool = False
94
+
95
+ class QueryResponse(BaseModel):
96
+ response: str
97
+
98
+ class InsertTextRequest(BaseModel):
99
+ text: str
100
+ description: Optional[str] = None
101
+
102
+ class InsertResponse(BaseModel):
103
+ status: str
104
+ message: str
105
+ document_count: int
106
+
107
+ async def get_embedding_dim(embedding_model: str) -> int:
108
+ """Get embedding dimensions for the specified model"""
109
+ test_text = ["This is a test sentence."]
110
+ embedding = await openai_embedding(test_text, model=embedding_model)
111
+ return embedding.shape[1]
112
+
113
+ def create_app(args):
114
+ # Setup logging
115
+ logging.basicConfig(format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level))
116
+
117
+ # Initialize FastAPI app
118
+ app = FastAPI(
119
+ title="LightRAG API",
120
+ description="API for querying text using LightRAG with OpenAI integration"
121
+ )
122
+
123
+ # Create working directory if it doesn't exist
124
+ Path(args.working_dir).mkdir(parents=True, exist_ok=True)
125
+
126
+ # Initialize document manager
127
+ doc_manager = DocumentManager(args.input_dir)
128
+
129
+ # Get embedding dimensions
130
+ embedding_dim = asyncio.run(get_embedding_dim(args.embedding_model))
131
+
132
+ # Initialize RAG with OpenAI configuration
133
+ rag = LightRAG(
134
+ working_dir=args.working_dir,
135
+ llm_model_func=async_openai_complete,
136
+ llm_model_name=args.model,
137
+ llm_model_max_token_size=args.max_tokens,
138
+ embedding_func=EmbeddingFunc(
139
+ embedding_dim=embedding_dim,
140
+ max_token_size=args.max_embed_tokens,
141
+ func=lambda texts: openai_embedding(texts, model=args.embedding_model),
142
+ ),
143
+ )
144
+
145
+ async def async_openai_complete(prompt, system_prompt=None, history_messages=[], **kwargs):
146
+ """Async wrapper for OpenAI completion"""
147
+ return await openai_complete_if_cache(
148
+ args.model,
149
+ prompt,
150
+ system_prompt=system_prompt,
151
+ history_messages=history_messages,
152
+ **kwargs
153
+ )
154
+ @app.on_event("startup")
155
+ async def startup_event():
156
+ """Index all files in input directory during startup"""
157
+ try:
158
+ new_files = doc_manager.scan_directory()
159
+ for file_path in new_files:
160
+ try:
161
+ # Use async file reading
162
+ async with aiofiles.open(file_path, 'r', encoding='utf-8') as f:
163
+ content = await f.read()
164
+ # Use the async version of insert directly
165
+ await rag.ainsert(content)
166
+ doc_manager.mark_as_indexed(file_path)
167
+ logging.info(f"Indexed file: {file_path}")
168
+ except Exception as e:
169
+ trace_exception(e)
170
+ logging.error(f"Error indexing file {file_path}: {str(e)}")
171
+
172
+ logging.info(f"Indexed {len(new_files)} documents from {args.input_dir}")
173
+
174
+ except Exception as e:
175
+ logging.error(f"Error during startup indexing: {str(e)}")
176
+
177
+ @app.post("/documents/scan")
178
+ async def scan_for_new_documents():
179
+ """Manually trigger scanning for new documents"""
180
+ try:
181
+ new_files = doc_manager.scan_directory()
182
+ indexed_count = 0
183
+
184
+ for file_path in new_files:
185
+ try:
186
+ with open(file_path, 'r', encoding='utf-8') as f:
187
+ content = f.read()
188
+ rag.insert(content)
189
+ doc_manager.mark_as_indexed(file_path)
190
+ indexed_count += 1
191
+ except Exception as e:
192
+ logging.error(f"Error indexing file {file_path}: {str(e)}")
193
+
194
+ return {
195
+ "status": "success",
196
+ "indexed_count": indexed_count,
197
+ "total_documents": len(doc_manager.indexed_files)
198
+ }
199
+ except Exception as e:
200
+ raise HTTPException(status_code=500, detail=str(e))
201
+
202
+ @app.post("/documents/upload")
203
+ async def upload_to_input_dir(file: UploadFile = File(...)):
204
+ """Upload a file to the input directory"""
205
+ try:
206
+ if not doc_manager.is_supported_file(file.filename):
207
+ raise HTTPException(
208
+ status_code=400,
209
+ detail=f"Unsupported file type. Supported types: {doc_manager.supported_extensions}"
210
+ )
211
+
212
+ file_path = doc_manager.input_dir / file.filename
213
+ with open(file_path, "wb") as buffer:
214
+ shutil.copyfileobj(file.file, buffer)
215
+
216
+ # Immediately index the uploaded file
217
+ with open(file_path, "r", encoding="utf-8") as f:
218
+ content = f.read()
219
+ rag.insert(content)
220
+ doc_manager.mark_as_indexed(file_path)
221
+
222
+ return {
223
+ "status": "success",
224
+ "message": f"File uploaded and indexed: {file.filename}",
225
+ "total_documents": len(doc_manager.indexed_files)
226
+ }
227
+ except Exception as e:
228
+ raise HTTPException(status_code=500, detail=str(e))
229
+
230
+ @app.post("/query", response_model=QueryResponse)
231
+ async def query_text(request: QueryRequest):
232
+ try:
233
+ response = await rag.aquery(
234
+ request.query,
235
+ param=QueryParam(mode=request.mode, stream=request.stream)
236
+ )
237
+
238
+ if request.stream:
239
+ result = ""
240
+ async for chunk in response:
241
+ result += chunk
242
+ return QueryResponse(response=result)
243
+ else:
244
+ return QueryResponse(response=response)
245
+ except Exception as e:
246
+ raise HTTPException(status_code=500, detail=str(e))
247
+
248
+ @app.post("/query/stream")
249
+ async def query_text_stream(request: QueryRequest):
250
+ try:
251
+ response = rag.query(
252
+ request.query,
253
+ param=QueryParam(mode=request.mode, stream=True)
254
+ )
255
+
256
+ async def stream_generator():
257
+ async for chunk in response:
258
+ yield chunk
259
+
260
+ return stream_generator()
261
+ except Exception as e:
262
+ raise HTTPException(status_code=500, detail=str(e))
263
+
264
+ @app.post("/documents/text", response_model=InsertResponse)
265
+ async def insert_text(request: InsertTextRequest):
266
+ try:
267
+ rag.insert(request.text)
268
+ return InsertResponse(
269
+ status="success",
270
+ message="Text successfully inserted",
271
+ document_count=len(rag)
272
+ )
273
+ except Exception as e:
274
+ raise HTTPException(status_code=500, detail=str(e))
275
+
276
+ @app.post("/documents/file", response_model=InsertResponse)
277
+ async def insert_file(
278
+ file: UploadFile = File(...),
279
+ description: str = Form(None)
280
+ ):
281
+ try:
282
+ content = await file.read()
283
+
284
+ if file.filename.endswith(('.txt', '.md')):
285
+ text = content.decode('utf-8')
286
+ rag.insert(text)
287
+ else:
288
+ raise HTTPException(
289
+ status_code=400,
290
+ detail="Unsupported file type. Only .txt and .md files are supported"
291
+ )
292
+
293
+ return InsertResponse(
294
+ status="success",
295
+ message=f"File '{file.filename}' successfully inserted",
296
+ document_count=len(rag)
297
+ )
298
+ except UnicodeDecodeError:
299
+ raise HTTPException(status_code=400, detail="File encoding not supported")
300
+ except Exception as e:
301
+ raise HTTPException(status_code=500, detail=str(e))
302
+
303
+ @app.post("/documents/batch", response_model=InsertResponse)
304
+ async def insert_batch(files: List[UploadFile] = File(...)):
305
+ try:
306
+ inserted_count = 0
307
+ failed_files = []
308
+
309
+ for file in files:
310
+ try:
311
+ content = await file.read()
312
+ if file.filename.endswith(('.txt', '.md')):
313
+ text = content.decode('utf-8')
314
+ rag.insert(text)
315
+ inserted_count += 1
316
+ else:
317
+ failed_files.append(f"{file.filename} (unsupported type)")
318
+ except Exception as e:
319
+ failed_files.append(f"{file.filename} ({str(e)})")
320
+
321
+ status_message = f"Successfully inserted {inserted_count} documents"
322
+ if failed_files:
323
+ status_message += f". Failed files: {', '.join(failed_files)}"
324
+
325
+ return InsertResponse(
326
+ status="success" if inserted_count > 0 else "partial_success",
327
+ message=status_message,
328
+ document_count=len(rag)
329
+ )
330
+ except Exception as e:
331
+ raise HTTPException(status_code=500, detail=str(e))
332
+
333
+ @app.delete("/documents", response_model=InsertResponse)
334
+ async def clear_documents():
335
+ try:
336
+ rag.text_chunks = []
337
+ rag.entities_vdb = None
338
+ rag.relationships_vdb = None
339
+ return InsertResponse(
340
+ status="success",
341
+ message="All documents cleared successfully",
342
+ document_count=0
343
+ )
344
+ except Exception as e:
345
+ raise HTTPException(status_code=500, detail=str(e))
346
+
347
+ @app.get("/health")
348
+ async def get_status():
349
+ """Get current system status"""
350
+ return {
351
+ "status": "healthy",
352
+ "working_directory": str(args.working_dir),
353
+ "input_directory": str(args.input_dir),
354
+ "indexed_files": len(doc_manager.indexed_files),
355
+ "configuration": {
356
+ "model": args.model,
357
+ "embedding_model": args.embedding_model,
358
+ "max_tokens": args.max_tokens,
359
+ "embedding_dim": embedding_dim
360
+ }
361
+ }
362
+
363
+ return app
364
+
365
+ if __name__ == "__main__":
366
+ args = parse_args()
367
+ import uvicorn
368
+ app = create_app(args)
369
+ uvicorn.run(app, host=args.host, port=args.port)