gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

updated

Browse Source
This commit is contained in:
Saifeddine ALOUI
2024-12-17 23:36:30 +01:00
parent b9c75dc2cd
commit 3d721e3db8
1 changed files with 129 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

178
api/ollama_lightrag_server.py
View File

@@ -10,52 +10,74 @@ from lightrag.llm import ollama_model_complete, ollama_embedding
from lightrag.utils import EmbeddingFunc
from typing import Optional, List
from enum import Enum
import io
from pathlib import Path
import shutil
def parse_args():
parser = argparse.ArgumentParser(
description="""
LightRAG FastAPI Server
======================
A REST API server for text querying using LightRAG. Supports multiple search modes,
streaming responses, and document management.
Features:
- Multiple search modes (naive, local, global, hybrid)
- Streaming and non-streaming responses
- Document insertion and management
- Configurable model parameters
- REST API with automatic documentation
""",
formatter_class=argparse.RawDescriptionHelpFormatter
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=8000, help='Server port (default: 8000)')
# 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='gemma2:2b', help='LLM model name (default: gemma2:2b)')
parser.add_argument('--embedding-model', default='nomic-embed-text', help='Embedding model name (default: nomic-embed-text)')
parser.add_argument('--ollama-host', default='http://localhost:11434', help='Ollama host URL (default: http://localhost:11434)')
parser.add_argument('--embedding-model', default='nomic-embed-text',
help='Embedding model name (default: nomic-embed-text)')
parser.add_argument('--ollama-host', default='http://localhost:11434',
help='Ollama host URL (default: http://localhost:11434)')
# RAG configuration
parser.add_argument('--working-dir', default='./dickens', help='Working directory for RAG (default: ./dickens)')
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=768, help='Embedding dimensions (default: 768)')
parser.add_argument('--max-embed-tokens', type=int, default=8192, help='Maximum embedding token size (default: 8192)')
# Input configuration
parser.add_argument('--input-file', default='./book.txt', help='Initial input file to process (default: ./book.txt)')
parser.add_argument('--embedding-dim', type=int, default=768,
help='Embedding dimensions (default: 768)')
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('--log-level', default='INFO',
choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
help='Logging level (default: INFO)')
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"
@@ -87,25 +109,14 @@ def create_app(args):
# Initialize FastAPI app
app = FastAPI(
title="LightRAG API",
description="""
API for querying text using LightRAG.
Configuration:
- Model: {model}
- Embedding Model: {embed_model}
- Working Directory: {work_dir}
- Max Tokens: {max_tokens}
""".format(
model=args.model,
embed_model=args.embedding_model,
work_dir=args.working_dir,
max_tokens=args.max_tokens
)
description="API for querying text using LightRAG with separate storage and input directories"
)
# Create working directory if it doesn't exist
if not os.path.exists(args.working_dir):
os.makedirs(args.working_dir)
Path(args.working_dir).mkdir(parents=True, exist_ok=True)
# Initialize document manager
doc_manager = DocumentManager(args.input_dir)
# Initialize RAG
rag = LightRAG(
@@ -126,11 +137,76 @@ def create_app(args):
@app.on_event("startup")
async def startup_event():
"""Index all files in input directory during startup"""
try:
with open(args.input_file, "r", encoding="utf-8") as f:
rag.insert(f.read())
except FileNotFoundError:
logging.warning(f"Input file {args.input_file} not found. Please ensure the file exists before querying.")
new_files = doc_manager.scan_directory()
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)
logging.info(f"Indexed file: {file_path}")
except Exception as 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")
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")
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)
async def query_text(request: QueryRequest):
@@ -249,14 +325,18 @@ def create_app(args):
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/health")
async def health_check():
@app.get("/status")
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,
"working_dir": args.working_dir,
"max_tokens": args.max_tokens,
"ollama_host": args.ollama_host
}