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

fixed linting

Browse Source
This commit is contained in:
Saifeddine ALOUI
2025-01-14 23:11:23 +01:00
parent c3aba5423f
commit 29661c92da
2 changed files with 50 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
lightrag/api/lightrag_server.py
View File

@@ -175,7 +175,11 @@ def parse_args():
class DocumentManager: class DocumentManager:
"""Handles document operations and tracking""" """Handles document operations and tracking"""
def __init__(self, input_dir: str, supported_extensions: tuple = (".txt", ".md", ".pdf", ".docx", ".pptx")): def __init__(
self,
input_dir: str,
supported_extensions: tuple = (".txt", ".md", ".pdf", ".docx", ".pptx"),
):
self.input_dir = Path(input_dir) self.input_dir = Path(input_dir)
self.supported_extensions = supported_extensions self.supported_extensions = supported_extensions
self.indexed_files = set() self.indexed_files = set()
@@ -357,26 +361,22 @@ def create_app(args):
), ),
) )
async def index_file(file_path: Union[str, Path]) -> None: async def index_file(file_path: Union[str, Path]) -> None:
""" Index all files inside the folder with support for multiple file formats """Index all files inside the folder with support for multiple file formats
Args: Args:
file_path: Path to the file to be indexed (str or Path object) file_path: Path to the file to be indexed (str or Path object)
Raises: Raises:
ValueError: If file format is not supported ValueError: If file format is not supported
FileNotFoundError: If file doesn't exist FileNotFoundError: If file doesn't exist
""" """
if not pm.is_installed("aiofiles"): if not pm.is_installed("aiofiles"):
pm.install("aiofiles") pm.install("aiofiles")
import aiofiles
# Convert to Path object if string # Convert to Path object if string
file_path = Path(file_path) file_path = Path(file_path)
# Check if file exists # Check if file exists
if not file_path.exists(): if not file_path.exists():
raise FileNotFoundError(f"File not found: {file_path}") raise FileNotFoundError(f"File not found: {file_path}")
@@ -384,23 +384,24 @@ def create_app(args):
content = "" content = ""
# Get file extension in lowercase # Get file extension in lowercase
ext = file_path.suffix.lower() ext = file_path.suffix.lower()
match ext: match ext:
case ".txt" | ".md": case ".txt" | ".md":
# Text files handling # Text files handling
async with aiofiles.open(file_path, "r", encoding="utf-8") as f: async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
content = await f.read() content = await f.read()
case ".pdf": case ".pdf":
if not pm.is_installed("pypdf2"): if not pm.is_installed("pypdf2"):
pm.install("pypdf2") pm.install("pypdf2")
from pypdf2 import PdfReader from pypdf2 import PdfReader
# PDF handling # PDF handling
reader = PdfReader(str(file_path)) reader = PdfReader(str(file_path))
content = "" content = ""
for page in reader.pages: for page in reader.pages:
content += page.extract_text() + "\n" content += page.extract_text() + "\n"
case ".docx": case ".docx":
if not pm.is_installed("docx"): if not pm.is_installed("docx"):
pm.install("docx") pm.install("docx")
@@ -409,11 +410,12 @@ def create_app(args):
# Word document handling # Word document handling
doc = Document(file_path) doc = Document(file_path)
content = "\n".join([paragraph.text for paragraph in doc.paragraphs]) content = "\n".join([paragraph.text for paragraph in doc.paragraphs])
case ".pptx": case ".pptx":
if not pm.is_installed("pptx"): if not pm.is_installed("pptx"):
pm.install("pptx") pm.install("pptx")
from pptx import Presentation from pptx import Presentation
# PowerPoint handling # PowerPoint handling
prs = Presentation(file_path) prs = Presentation(file_path)
content = "" content = ""
@@ -421,7 +423,7 @@ def create_app(args):
for shape in slide.shapes: for shape in slide.shapes:
if hasattr(shape, "text"): if hasattr(shape, "text"):
content += shape.text + "\n" content += shape.text + "\n"
case _: case _:
raise ValueError(f"Unsupported file format: {ext}") raise ValueError(f"Unsupported file format: {ext}")
@@ -433,9 +435,6 @@ def create_app(args):
else: else:
logging.warning(f"No content extracted from file: {file_path}") logging.warning(f"No content extracted from file: {file_path}")
@app.on_event("startup") @app.on_event("startup")
async def startup_event(): async def startup_event():
"""Index all files in input directory during startup""" """Index all files in input directory during startup"""
@@ -559,6 +558,7 @@ def create_app(args):
) )
except Exception as e: except Exception as e:
raise HTTPException(status_code=500, detail=str(e)) raise HTTPException(status_code=500, detail=str(e))
@app.post( @app.post(
"/documents/file", "/documents/file",
response_model=InsertResponse, response_model=InsertResponse,
@@ -566,14 +566,14 @@ def create_app(args):
) )
async def insert_file(file: UploadFile = File(...), description: str = Form(None)): async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
"""Insert a file directly into the RAG system """Insert a file directly into the RAG system
Args: Args:
file: Uploaded file file: Uploaded file
description: Optional description of the file description: Optional description of the file
Returns: Returns:
InsertResponse: Status of the insertion operation InsertResponse: Status of the insertion operation
Raises: Raises:
HTTPException: For unsupported file types or processing errors HTTPException: For unsupported file types or processing errors
""" """
@@ -581,19 +581,19 @@ def create_app(args):
content = "" content = ""
# Get file extension in lowercase # Get file extension in lowercase
ext = Path(file.filename).suffix.lower() ext = Path(file.filename).suffix.lower()
match ext: match ext:
case ".txt" | ".md": case ".txt" | ".md":
# Text files handling # Text files handling
text_content = await file.read() text_content = await file.read()
content = text_content.decode("utf-8") content = text_content.decode("utf-8")
case ".pdf": case ".pdf":
if not pm.is_installed("pypdf2"): if not pm.is_installed("pypdf2"):
pm.install("pypdf2") pm.install("pypdf2")
from pypdf2 import PdfReader from pypdf2 import PdfReader
from io import BytesIO from io import BytesIO
# Read PDF from memory # Read PDF from memory
pdf_content = await file.read() pdf_content = await file.read()
pdf_file = BytesIO(pdf_content) pdf_file = BytesIO(pdf_content)
@@ -601,25 +601,27 @@ def create_app(args):
content = "" content = ""
for page in reader.pages: for page in reader.pages:
content += page.extract_text() + "\n" content += page.extract_text() + "\n"
case ".docx": case ".docx":
if not pm.is_installed("docx"): if not pm.is_installed("docx"):
pm.install("docx") pm.install("docx")
from docx import Document from docx import Document
from io import BytesIO from io import BytesIO
# Read DOCX from memory # Read DOCX from memory
docx_content = await file.read() docx_content = await file.read()
docx_file = BytesIO(docx_content) docx_file = BytesIO(docx_content)
doc = Document(docx_file) doc = Document(docx_file)
content = "\n".join([paragraph.text for paragraph in doc.paragraphs]) content = "\n".join(
[paragraph.text for paragraph in doc.paragraphs]
)
case ".pptx": case ".pptx":
if not pm.is_installed("pptx"): if not pm.is_installed("pptx"):
pm.install("pptx") pm.install("pptx")
from pptx import Presentation from pptx import Presentation
from io import BytesIO from io import BytesIO
# Read PPTX from memory # Read PPTX from memory
pptx_content = await file.read() pptx_content = await file.read()
pptx_file = BytesIO(pptx_content) pptx_file = BytesIO(pptx_content)
@@ -629,7 +631,7 @@ def create_app(args):
for shape in slide.shapes: for shape in slide.shapes:
if hasattr(shape, "text"): if hasattr(shape, "text"):
content += shape.text + "\n" content += shape.text + "\n"
case _: case _:
raise HTTPException( raise HTTPException(
status_code=400, status_code=400,
@@ -641,10 +643,10 @@ def create_app(args):
# Add description if provided # Add description if provided
if description: if description:
content = f"{description}\n\n{content}" content = f"{description}\n\n{content}"
await rag.ainsert(content) await rag.ainsert(content)
logging.info(f"Successfully indexed file: {file.filename}") logging.info(f"Successfully indexed file: {file.filename}")
return InsertResponse( return InsertResponse(
status="success", status="success",
message=f"File '{file.filename}' successfully inserted", message=f"File '{file.filename}' successfully inserted",
@@ -661,6 +663,7 @@ def create_app(args):
except Exception as e: except Exception as e:
logging.error(f"Error processing file {file.filename}: {str(e)}") logging.error(f"Error processing file {file.filename}: {str(e)}")
raise HTTPException(status_code=500, detail=str(e)) raise HTTPException(status_code=500, detail=str(e))
@app.post( @app.post(
"/documents/batch", "/documents/batch",
response_model=InsertResponse, response_model=InsertResponse,
@@ -668,13 +671,13 @@ def create_app(args):
) )
async def insert_batch(files: List[UploadFile] = File(...)): async def insert_batch(files: List[UploadFile] = File(...)):
"""Process multiple files in batch mode """Process multiple files in batch mode
Args: Args:
files: List of files to process files: List of files to process
Returns: Returns:
InsertResponse: Status of the batch insertion operation InsertResponse: Status of the batch insertion operation
Raises: Raises:
HTTPException: For processing errors HTTPException: For processing errors
""" """
@@ -686,41 +689,43 @@ def create_app(args):
try: try:
content = "" content = ""
ext = Path(file.filename).suffix.lower() ext = Path(file.filename).suffix.lower()
match ext: match ext:
case ".txt" | ".md": case ".txt" | ".md":
text_content = await file.read() text_content = await file.read()
content = text_content.decode("utf-8") content = text_content.decode("utf-8")
case ".pdf": case ".pdf":
if not pm.is_installed("pypdf2"): if not pm.is_installed("pypdf2"):
pm.install("pypdf2") pm.install("pypdf2")
from pypdf2 import PdfReader from pypdf2 import PdfReader
from io import BytesIO from io import BytesIO
pdf_content = await file.read() pdf_content = await file.read()
pdf_file = BytesIO(pdf_content) pdf_file = BytesIO(pdf_content)
reader = PdfReader(pdf_file) reader = PdfReader(pdf_file)
for page in reader.pages: for page in reader.pages:
content += page.extract_text() + "\n" content += page.extract_text() + "\n"
case ".docx": case ".docx":
if not pm.is_installed("docx"): if not pm.is_installed("docx"):
pm.install("docx") pm.install("docx")
from docx import Document from docx import Document
from io import BytesIO from io import BytesIO
docx_content = await file.read() docx_content = await file.read()
docx_file = BytesIO(docx_content) docx_file = BytesIO(docx_content)
doc = Document(docx_file) doc = Document(docx_file)
content = "\n".join([paragraph.text for paragraph in doc.paragraphs]) content = "\n".join(
[paragraph.text for paragraph in doc.paragraphs]
)
case ".pptx": case ".pptx":
if not pm.is_installed("pptx"): if not pm.is_installed("pptx"):
pm.install("pptx") pm.install("pptx")
from pptx import Presentation from pptx import Presentation
from io import BytesIO from io import BytesIO
pptx_content = await file.read() pptx_content = await file.read()
pptx_file = BytesIO(pptx_content) pptx_file = BytesIO(pptx_content)
prs = Presentation(pptx_file) prs = Presentation(pptx_file)
@@ -728,7 +733,7 @@ def create_app(args):
for shape in slide.shapes: for shape in slide.shapes:
if hasattr(shape, "text"): if hasattr(shape, "text"):
content += shape.text + "\n" content += shape.text + "\n"
case _: case _:
failed_files.append(f"{file.filename} (unsupported type)") failed_files.append(f"{file.filename} (unsupported type)")
continue continue
@@ -771,7 +776,6 @@ def create_app(args):
logging.error(f"Batch processing error: {str(e)}") logging.error(f"Batch processing error: {str(e)}")
raise HTTPException(status_code=500, detail=str(e)) raise HTTPException(status_code=500, detail=str(e))
@app.delete( @app.delete(
"/documents", "/documents",
response_model=InsertResponse, response_model=InsertResponse,

2
lightrag/api/requirements.txt
View File

@@ -7,6 +7,7 @@ nest_asyncio
numpy numpy
ollama ollama
openai openai
pipmaster
python-dotenv python-dotenv
python-multipart python-multipart
tenacity tenacity
@@ -15,4 +16,3 @@ torch
tqdm tqdm
transformers transformers
uvicorn uvicorn
pipmaster