diff --git a/api/ollama_lightrag_server.py b/api/ollama_lightrag_server.py index 383e04cb..589cda0e 100644 --- a/api/ollama_lightrag_server.py +++ b/api/ollama_lightrag_server.py @@ -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 }