Added openai api
Browse files- api/{README.md → README_OLLAMA.md} +0 -0
- api/README_OPENAI.md +172 -0
- api/openai_lightrag_server.py +369 -0
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)
|