diff --git a/README.md b/README.md
index ad405e90..dd215b04 100644
--- a/README.md
+++ b/README.md
@@ -465,7 +465,36 @@ For production level scenarios you will most likely want to leverage an enterpri
>
> You can Compile the AGE from source code and fix it.
+### Using Faiss for Storage
+- Install the required dependencies:
+```
+pip install faiss-cpu
+```
+You can also install `faiss-gpu` if you have GPU support.
+- Here we are using `sentence-transformers` but you can also use `OpenAIEmbedding` model with `3072` dimensions.
+
+```
+async def embedding_func(texts: list[str]) -> np.ndarray:
+ model = SentenceTransformer('all-MiniLM-L6-v2')
+ embeddings = model.encode(texts, convert_to_numpy=True)
+ return embeddings
+
+# Initialize LightRAG with the LLM model function and embedding function
+ rag = LightRAG(
+ working_dir=WORKING_DIR,
+ llm_model_func=llm_model_func,
+ embedding_func=EmbeddingFunc(
+ embedding_dim=384,
+ max_token_size=8192,
+ func=embedding_func,
+ ),
+ vector_storage="FaissVectorDBStorage",
+ vector_db_storage_cls_kwargs={
+ "cosine_better_than_threshold": 0.3 # Your desired threshold
+ }
+ )
+```
### Insert Custom KG
diff --git a/examples/test_faiss.py b/examples/test_faiss.py
new file mode 100644
index 00000000..ab0ef9f7
--- /dev/null
+++ b/examples/test_faiss.py
@@ -0,0 +1,99 @@
+import os
+import logging
+import numpy as np
+
+from dotenv import load_dotenv
+from sentence_transformers import SentenceTransformer
+
+from openai import AzureOpenAI
+from lightrag import LightRAG, QueryParam
+from lightrag.utils import EmbeddingFunc
+
+# Configure Logging
+logging.basicConfig(level=logging.INFO)
+
+# Load environment variables from .env file
+load_dotenv()
+AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION")
+AZURE_OPENAI_DEPLOYMENT = os.getenv("AZURE_OPENAI_DEPLOYMENT")
+AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
+AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
+
+
+async def llm_model_func(
+ prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+ # Create a client for AzureOpenAI
+ client = AzureOpenAI(
+ api_key=AZURE_OPENAI_API_KEY,
+ api_version=AZURE_OPENAI_API_VERSION,
+ azure_endpoint=AZURE_OPENAI_ENDPOINT,
+ )
+
+ # Build the messages list for the conversation
+ messages = []
+ if system_prompt:
+ messages.append({"role": "system", "content": system_prompt})
+ if history_messages:
+ messages.extend(history_messages)
+ messages.append({"role": "user", "content": prompt})
+
+ # Call the LLM
+ chat_completion = client.chat.completions.create(
+ model=AZURE_OPENAI_DEPLOYMENT,
+ messages=messages,
+ temperature=kwargs.get("temperature", 0),
+ top_p=kwargs.get("top_p", 1),
+ n=kwargs.get("n", 1),
+ )
+
+ return chat_completion.choices[0].message.content
+
+
+async def embedding_func(texts: list[str]) -> np.ndarray:
+ model = SentenceTransformer("all-MiniLM-L6-v2")
+ embeddings = model.encode(texts, convert_to_numpy=True)
+ return embeddings
+
+
+def main():
+ WORKING_DIR = "./dickens"
+
+ # Initialize LightRAG with the LLM model function and embedding function
+ rag = LightRAG(
+ working_dir=WORKING_DIR,
+ llm_model_func=llm_model_func,
+ embedding_func=EmbeddingFunc(
+ embedding_dim=384,
+ max_token_size=8192,
+ func=embedding_func,
+ ),
+ vector_storage="FaissVectorDBStorage",
+ vector_db_storage_cls_kwargs={
+ "cosine_better_than_threshold": 0.3 # Your desired threshold
+ },
+ )
+
+ # Insert the custom chunks into LightRAG
+ book1 = open("./book_1.txt", encoding="utf-8")
+ book2 = open("./book_2.txt", encoding="utf-8")
+
+ rag.insert([book1.read(), book2.read()])
+
+ query_text = "What are the main themes?"
+
+ print("Result (Naive):")
+ print(rag.query(query_text, param=QueryParam(mode="naive")))
+
+ print("\nResult (Local):")
+ print(rag.query(query_text, param=QueryParam(mode="local")))
+
+ print("\nResult (Global):")
+ print(rag.query(query_text, param=QueryParam(mode="global")))
+
+ print("\nResult (Hybrid):")
+ print(rag.query(query_text, param=QueryParam(mode="hybrid")))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py
index e162f5ec..e1b24731 100644
--- a/lightrag/api/lightrag_server.py
+++ b/lightrag/api/lightrag_server.py
@@ -1,9 +1,37 @@
-from fastapi import FastAPI, HTTPException, File, UploadFile, Form, Request
+from fastapi import (
+ FastAPI,
+ HTTPException,
+ File,
+ UploadFile,
+ Form,
+ Request,
+ BackgroundTasks,
+)
+
+# Backend (Python)
+# Add this to store progress globally
+from typing import Dict
+import threading
+
+# Global progress tracker
+scan_progress: Dict = {
+ "is_scanning": False,
+ "current_file": "",
+ "indexed_count": 0,
+ "total_files": 0,
+ "progress": 0,
+}
+
+# Lock for thread-safe operations
+progress_lock = threading.Lock()
+
+import json
+import os
+
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import logging
import argparse
-import json
import time
import re
from typing import List, Dict, Any, Optional, Union
@@ -16,7 +44,6 @@ from pathlib import Path
import shutil
import aiofiles
from ascii_colors import trace_exception, ASCIIColors
-import os
import sys
import configparser
@@ -538,7 +565,7 @@ class DocumentManager:
# Create input directory if it doesn't exist
self.input_dir.mkdir(parents=True, exist_ok=True)
- def scan_directory(self) -> List[Path]:
+ def scan_directory_for_new_files(self) -> List[Path]:
"""Scan input directory for new files"""
new_files = []
for ext in self.supported_extensions:
@@ -547,6 +574,14 @@ class DocumentManager:
new_files.append(file_path)
return new_files
+ 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}"):
+ 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)
@@ -730,7 +765,7 @@ def create_app(args):
# Startup logic
if args.auto_scan_at_startup:
try:
- new_files = doc_manager.scan_directory()
+ new_files = doc_manager.scan_directory_for_new_files()
for file_path in new_files:
try:
await index_file(file_path)
@@ -983,42 +1018,59 @@ def create_app(args):
logging.warning(f"No content extracted from file: {file_path}")
@app.post("/documents/scan", dependencies=[Depends(optional_api_key)])
- async def scan_for_new_documents():
- """
- Manually trigger scanning for new documents in the directory managed by `doc_manager`.
+ async def scan_for_new_documents(background_tasks: BackgroundTasks):
+ """Trigger the scanning process"""
+ global scan_progress
- This endpoint facilitates manual initiation of a document scan to identify and index new files.
- It processes all newly detected files, attempts indexing each file, logs any errors that occur,
- and returns a summary of the operation.
+ with progress_lock:
+ if scan_progress["is_scanning"]:
+ return {"status": "already_scanning"}
- Returns:
- dict: A dictionary containing:
- - "status" (str): Indicates success or failure of the scanning process.
- - "indexed_count" (int): The number of successfully indexed documents.
- - "total_documents" (int): Total number of documents that have been indexed so far.
+ scan_progress["is_scanning"] = True
+ scan_progress["indexed_count"] = 0
+ scan_progress["progress"] = 0
+
+ # Start the scanning process in the background
+ background_tasks.add_task(run_scanning_process)
+
+ return {"status": "scanning_started"}
+
+ async def run_scanning_process():
+ """Background task to scan and index documents"""
+ global scan_progress
- Raises:
- HTTPException: If an error occurs during the document scanning process, a 500 status
- code is returned with details about the exception.
- """
try:
- new_files = doc_manager.scan_directory()
- indexed_count = 0
+ new_files = doc_manager.scan_directory_for_new_files()
+ scan_progress["total_files"] = len(new_files)
for file_path in new_files:
try:
+ with progress_lock:
+ scan_progress["current_file"] = os.path.basename(file_path)
+
await index_file(file_path)
- indexed_count += 1
+
+ with progress_lock:
+ scan_progress["indexed_count"] += 1
+ scan_progress["progress"] = (
+ scan_progress["indexed_count"]
+ / scan_progress["total_files"]
+ ) * 100
+
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))
+ logging.error(f"Error during scanning process: {str(e)}")
+ finally:
+ with progress_lock:
+ scan_progress["is_scanning"] = False
+
+ @app.get("/documents/scan-progress")
+ async def get_scan_progress():
+ """Get the current scanning progress"""
+ with progress_lock:
+ return scan_progress
@app.post("/documents/upload", dependencies=[Depends(optional_api_key)])
async def upload_to_input_dir(file: UploadFile = File(...)):
@@ -1849,7 +1901,7 @@ def create_app(args):
"status": "healthy",
"working_directory": str(args.working_dir),
"input_directory": str(args.input_dir),
- "indexed_files": files,
+ "indexed_files": [str(f) for f in files],
"indexed_files_count": len(files),
"configuration": {
# LLM configuration binding/host address (if applicable)/model (if applicable)
diff --git a/lightrag/api/static/index.html b/lightrag/api/static/index.html
index 60900c03..c9659d5e 100644
--- a/lightrag/api/static/index.html
+++ b/lightrag/api/static/index.html
@@ -98,7 +98,7 @@
-
+