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

Fix linting

Browse Source
This commit is contained in:
yangdx
2025-02-18 16:33:57 +08:00
parent 75ee4592b8
commit 5acd054180
1 changed files with 15 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
lightrag/api/lightrag_server.py
View File

@@ -41,6 +41,8 @@ from .ollama_api import (
OllamaAPI, OllamaAPI,
) )
from .ollama_api import ollama_server_infos from .ollama_api import ollama_server_infos
def get_db_type_from_storage_class(class_name: str) -> str | None: def get_db_type_from_storage_class(class_name: str) -> str | None:
"""Determine database type based on storage class name""" """Determine database type based on storage class name"""
if class_name.startswith("PG"): if class_name.startswith("PG"):
@@ -51,19 +53,24 @@ def get_db_type_from_storage_class(class_name: str) -> str | None:
return "tidb" return "tidb"
return None return None
def import_db_module(db_type: str): def import_db_module(db_type: str):
"""Dynamically import database module""" """Dynamically import database module"""
if db_type == "postgres": if db_type == "postgres":
from ..kg.postgres_impl import PostgreSQLDB from ..kg.postgres_impl import PostgreSQLDB
return PostgreSQLDB return PostgreSQLDB
elif db_type == "oracle": elif db_type == "oracle":
from ..kg.oracle_impl import OracleDB from ..kg.oracle_impl import OracleDB
return OracleDB return OracleDB
elif db_type == "tidb": elif db_type == "tidb":
from ..kg.tidb_impl import TiDB from ..kg.tidb_impl import TiDB
return TiDB return TiDB
return None return None
# Load environment variables # Load environment variables
try: try:
load_dotenv(override=True) load_dotenv(override=True)
@@ -901,7 +908,9 @@ def create_app(args):
# Check which database types are used # Check which database types are used
db_types = set() db_types = set()
for storage_name, storage_instance in storage_instances: for storage_name, storage_instance in storage_instances:
db_type = get_db_type_from_storage_class(storage_instance.__class__.__name__) db_type = get_db_type_from_storage_class(
storage_instance.__class__.__name__
)
if db_type: if db_type:
db_types.add(db_type) db_types.add(db_type)
@@ -926,7 +935,9 @@ def create_app(args):
# Inject database instances into storage classes # Inject database instances into storage classes
for storage_name, storage_instance in storage_instances: for storage_name, storage_instance in storage_instances:
db_type = get_db_type_from_storage_class(storage_instance.__class__.__name__) db_type = get_db_type_from_storage_class(
storage_instance.__class__.__name__
)
if db_type: if db_type:
if db_type not in db_instances: if db_type not in db_instances:
error_msg = f"Database type '{db_type}' is required by {storage_name} but not initialized" error_msg = f"Database type '{db_type}' is required by {storage_name} but not initialized"
@@ -966,7 +977,7 @@ def create_app(args):
db_names = { db_names = {
"postgres": "PostgreSQL", "postgres": "PostgreSQL",
"oracle": "Oracle", "oracle": "Oracle",
"tidb": "TiDB" "tidb": "TiDB",
} }
db_name = db_names.get(db_type, db_type) db_name = db_names.get(db_type, db_type)
logger.info(f"Closed {db_name} database connection pool") logger.info(f"Closed {db_name} database connection pool")
@@ -1289,7 +1300,7 @@ def create_app(args):
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 # type: ignore from PyPDF2 import PdfReader # type: ignore
from io import BytesIO from io import BytesIO
pdf_file = BytesIO(file) pdf_file = BytesIO(file)