updated

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="""
+        description="LightRAG FastAPI Server with separate working and input directories"
-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
     )
 
     # 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('--embedding-model', default='nomic-embed-text', 
-    parser.add_argument('--ollama-host', default='http://localhost:11434', help='Ollama host URL (default: http://localhost:11434)')
+                       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('--embedding-dim', type=int, default=768, 
-    parser.add_argument('--max-embed-tokens', type=int, default=8192, help='Maximum embedding token size (default: 8192)')
+                       help='Embedding dimensions (default: 768)')
-    
+    parser.add_argument('--max-embed-tokens', type=int, default=8192, 
-    # Input configuration
+                       help='Maximum embedding token size (default: 8192)')
-    parser.add_argument('--input-file', default='./book.txt', help='Initial input file to process (default: ./book.txt)')
     
     # Logging configuration
-    parser.add_argument('--log-level', default='INFO', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
+    parser.add_argument('--log-level', default='INFO', 
-                      help='Logging 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="""
+        description="API for querying text using LightRAG with separate storage and input directories"
-        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
-        )
     )
 
     # Create working directory if it doesn't exist
-    if not os.path.exists(args.working_dir):
+    Path(args.working_dir).mkdir(parents=True, exist_ok=True)
-        os.makedirs(args.working_dir)
+
+    # 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:
+            new_files = doc_manager.scan_directory()
-                rag.insert(f.read())
+            for file_path in new_files:
-        except FileNotFoundError:
+                try:
-            logging.warning(f"Input file {args.input_file} not found. Please ensure the file exists before querying.")
+                    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
             }