[feat] Add API server implementation and endpoints
This commit is contained in:
119
README.md
119
README.md
@@ -397,6 +397,125 @@ if __name__ == "__main__":
|
||||
|
||||
</details>
|
||||
|
||||
## API Server Implementation
|
||||
|
||||
LightRAG also provides a FastAPI-based server implementation for RESTful API access to RAG operations. This allows you to run LightRAG as a service and interact with it through HTTP requests.
|
||||
|
||||
### Setting up the API Server
|
||||
<details>
|
||||
<summary>Click to expand setup instructions</summary>
|
||||
|
||||
1. First, ensure you have the required dependencies:
|
||||
```bash
|
||||
pip install fastapi uvicorn pydantic
|
||||
```
|
||||
|
||||
2. Set up your environment variables:
|
||||
```bash
|
||||
export RAG_DIR="your_index_directory" # Optional: Defaults to "index_default"
|
||||
```
|
||||
|
||||
3. Run the API server:
|
||||
```bash
|
||||
python examples/lightrag_api_openai_compatible_demo.py
|
||||
```
|
||||
|
||||
The server will start on `http://0.0.0.0:8020`.
|
||||
</details>
|
||||
|
||||
### API Endpoints
|
||||
|
||||
The API server provides the following endpoints:
|
||||
|
||||
#### 1. Query Endpoint
|
||||
<details>
|
||||
<summary>Click to view Query endpoint details</summary>
|
||||
|
||||
- **URL:** `/query`
|
||||
- **Method:** POST
|
||||
- **Body:**
|
||||
```json
|
||||
{
|
||||
"query": "Your question here",
|
||||
"mode": "hybrid" // Can be "naive", "local", "global", or "hybrid"
|
||||
}
|
||||
```
|
||||
- **Example:**
|
||||
```bash
|
||||
curl -X POST "http://127.0.0.1:8020/query" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"query": "What are the main themes?", "mode": "hybrid"}'
|
||||
```
|
||||
</details>
|
||||
|
||||
#### 2. Insert Text Endpoint
|
||||
<details>
|
||||
<summary>Click to view Insert Text endpoint details</summary>
|
||||
|
||||
- **URL:** `/insert`
|
||||
- **Method:** POST
|
||||
- **Body:**
|
||||
```json
|
||||
{
|
||||
"text": "Your text content here"
|
||||
}
|
||||
```
|
||||
- **Example:**
|
||||
```bash
|
||||
curl -X POST "http://127.0.0.1:8020/insert" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"text": "Content to be inserted into RAG"}'
|
||||
```
|
||||
</details>
|
||||
|
||||
#### 3. Insert File Endpoint
|
||||
<details>
|
||||
<summary>Click to view Insert File endpoint details</summary>
|
||||
|
||||
- **URL:** `/insert_file`
|
||||
- **Method:** POST
|
||||
- **Body:**
|
||||
```json
|
||||
{
|
||||
"file_path": "path/to/your/file.txt"
|
||||
}
|
||||
```
|
||||
- **Example:**
|
||||
```bash
|
||||
curl -X POST "http://127.0.0.1:8020/insert_file" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"file_path": "./book.txt"}'
|
||||
```
|
||||
</details>
|
||||
|
||||
#### 4. Health Check Endpoint
|
||||
<details>
|
||||
<summary>Click to view Health Check endpoint details</summary>
|
||||
|
||||
- **URL:** `/health`
|
||||
- **Method:** GET
|
||||
- **Example:**
|
||||
```bash
|
||||
curl -X GET "http://127.0.0.1:8020/health"
|
||||
```
|
||||
</details>
|
||||
|
||||
### Configuration
|
||||
|
||||
The API server can be configured using environment variables:
|
||||
- `RAG_DIR`: Directory for storing the RAG index (default: "index_default")
|
||||
- API keys and base URLs should be configured in the code for your specific LLM and embedding model providers
|
||||
|
||||
### Error Handling
|
||||
<details>
|
||||
<summary>Click to view error handling details</summary>
|
||||
|
||||
The API includes comprehensive error handling:
|
||||
- File not found errors (404)
|
||||
- Processing errors (500)
|
||||
- Supports multiple file encodings (UTF-8 and GBK)
|
||||
</details>
|
||||
|
||||
## Evaluation
|
||||
### Dataset
|
||||
The dataset used in LightRAG can be downloaded from [TommyChien/UltraDomain](https://huggingface.co/datasets/TommyChien/UltraDomain).
|
||||
|
153
examples/lightrag_api_openai_compatible_demo.py
Normal file
153
examples/lightrag_api_openai_compatible_demo.py
Normal file
@@ -0,0 +1,153 @@
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
import os
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm import openai_complete_if_cache, openai_embedding
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
import numpy as np
|
||||
from typing import Optional
|
||||
import asyncio
|
||||
import nest_asyncio
|
||||
|
||||
# Apply nest_asyncio to solve event loop issues
|
||||
nest_asyncio.apply()
|
||||
|
||||
DEFAULT_RAG_DIR="index_default"
|
||||
app = FastAPI(title="LightRAG API", description="API for RAG operations")
|
||||
|
||||
# Configure working directory
|
||||
WORKING_DIR = os.environ.get('RAG_DIR', f'{DEFAULT_RAG_DIR}')
|
||||
print(f"WORKING_DIR: {WORKING_DIR}")
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
# LLM model function
|
||||
async def llm_model_func(
|
||||
prompt, system_prompt=None, history_messages=[], **kwargs
|
||||
) -> str:
|
||||
return await openai_complete_if_cache(
|
||||
"gpt-4o-mini",
|
||||
prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages,
|
||||
api_key='YOUR_API_KEY',
|
||||
base_url="YourURL/v1",
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Embedding function
|
||||
async def embedding_func(texts: list[str]) -> np.ndarray:
|
||||
return await openai_embedding(
|
||||
texts,
|
||||
model="text-embedding-3-large",
|
||||
api_key='YOUR_API_KEY',
|
||||
base_url="YourURL/v1",
|
||||
)
|
||||
|
||||
# Initialize RAG instance
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=llm_model_func,
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=3072, max_token_size=8192, func=embedding_func
|
||||
),
|
||||
)
|
||||
|
||||
# Data models
|
||||
class QueryRequest(BaseModel):
|
||||
query: str
|
||||
mode: str = "hybrid"
|
||||
|
||||
class InsertRequest(BaseModel):
|
||||
text: str
|
||||
|
||||
class InsertFileRequest(BaseModel):
|
||||
file_path: str
|
||||
|
||||
class Response(BaseModel):
|
||||
status: str
|
||||
data: Optional[str] = None
|
||||
message: Optional[str] = None
|
||||
|
||||
# API routes
|
||||
@app.post("/query", response_model=Response)
|
||||
async def query_endpoint(request: QueryRequest):
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: rag.query(request.query, param=QueryParam(mode=request.mode))
|
||||
)
|
||||
return Response(
|
||||
status="success",
|
||||
data=result
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post("/insert", response_model=Response)
|
||||
async def insert_endpoint(request: InsertRequest):
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, lambda: rag.insert(request.text))
|
||||
return Response(
|
||||
status="success",
|
||||
message="Text inserted successfully"
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post("/insert_file", response_model=Response)
|
||||
async def insert_file(request: InsertFileRequest):
|
||||
try:
|
||||
# Check if file exists
|
||||
if not os.path.exists(request.file_path):
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"File not found: {request.file_path}"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(request.file_path, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
except UnicodeDecodeError:
|
||||
# If UTF-8 decoding fails, try other encodings
|
||||
with open(request.file_path, 'r', encoding='gbk') as f:
|
||||
content = f.read()
|
||||
|
||||
# Insert file content
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, lambda: rag.insert(content))
|
||||
|
||||
return Response(
|
||||
status="success",
|
||||
message=f"File content from {request.file_path} inserted successfully"
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
return {"status": "healthy"}
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=8020)
|
||||
|
||||
# Usage example
|
||||
# To run the server, use the following command in your terminal:
|
||||
# python lightrag_api_openai_compatible_demo.py
|
||||
|
||||
# Example requests:
|
||||
# 1. Query:
|
||||
# curl -X POST "http://127.0.0.1:8020/query" -H "Content-Type: application/json" -d '{"query": "your query here", "mode": "hybrid"}'
|
||||
|
||||
# 2. Insert text:
|
||||
# curl -X POST "http://127.0.0.1:8020/insert" -H "Content-Type: application/json" -d '{"text": "your text here"}'
|
||||
|
||||
# 3. Insert file:
|
||||
# curl -X POST "http://127.0.0.1:8020/insert_file" -H "Content-Type: application/json" -d '{"file_path": "path/to/your/file.txt"}'
|
||||
|
||||
# 4. Health check:
|
||||
# curl -X GET "http://127.0.0.1:8020/health"
|
Reference in New Issue
Block a user