yangdx commited on
Commit
395a6bb
·
1 Parent(s): a7c0d7b

Fix linting

Browse files
Files changed (1) hide show
  1. lightrag/api/lightrag_server.py +15 -4
lightrag/api/lightrag_server.py CHANGED
@@ -41,6 +41,8 @@ from .ollama_api import (
41
  OllamaAPI,
42
  )
43
  from .ollama_api import ollama_server_infos
 
 
44
  def get_db_type_from_storage_class(class_name: str) -> str | None:
45
  """Determine database type based on storage class name"""
46
  if class_name.startswith("PG"):
@@ -51,19 +53,24 @@ def get_db_type_from_storage_class(class_name: str) -> str | None:
51
  return "tidb"
52
  return None
53
 
 
54
  def import_db_module(db_type: str):
55
  """Dynamically import database module"""
56
  if db_type == "postgres":
57
  from ..kg.postgres_impl import PostgreSQLDB
 
58
  return PostgreSQLDB
59
  elif db_type == "oracle":
60
  from ..kg.oracle_impl import OracleDB
 
61
  return OracleDB
62
  elif db_type == "tidb":
63
  from ..kg.tidb_impl import TiDB
 
64
  return TiDB
65
  return None
66
 
 
67
  # Load environment variables
68
  try:
69
  load_dotenv(override=True)
@@ -901,7 +908,9 @@ def create_app(args):
901
  # Check which database types are used
902
  db_types = set()
903
  for storage_name, storage_instance in storage_instances:
904
- db_type = get_db_type_from_storage_class(storage_instance.__class__.__name__)
 
 
905
  if db_type:
906
  db_types.add(db_type)
907
 
@@ -926,7 +935,9 @@ def create_app(args):
926
 
927
  # Inject database instances into storage classes
928
  for storage_name, storage_instance in storage_instances:
929
- db_type = get_db_type_from_storage_class(storage_instance.__class__.__name__)
 
 
930
  if db_type:
931
  if db_type not in db_instances:
932
  error_msg = f"Database type '{db_type}' is required by {storage_name} but not initialized"
@@ -966,7 +977,7 @@ def create_app(args):
966
  db_names = {
967
  "postgres": "PostgreSQL",
968
  "oracle": "Oracle",
969
- "tidb": "TiDB"
970
  }
971
  db_name = db_names.get(db_type, db_type)
972
  logger.info(f"Closed {db_name} database connection pool")
@@ -1289,7 +1300,7 @@ def create_app(args):
1289
  case ".pdf":
1290
  if not pm.is_installed("pypdf2"):
1291
  pm.install("pypdf2")
1292
- from PyPDF2 import PdfReader # type: ignore
1293
  from io import BytesIO
1294
 
1295
  pdf_file = BytesIO(file)
 
41
  OllamaAPI,
42
  )
43
  from .ollama_api import ollama_server_infos
44
+
45
+
46
  def get_db_type_from_storage_class(class_name: str) -> str | None:
47
  """Determine database type based on storage class name"""
48
  if class_name.startswith("PG"):
 
53
  return "tidb"
54
  return None
55
 
56
+
57
  def import_db_module(db_type: str):
58
  """Dynamically import database module"""
59
  if db_type == "postgres":
60
  from ..kg.postgres_impl import PostgreSQLDB
61
+
62
  return PostgreSQLDB
63
  elif db_type == "oracle":
64
  from ..kg.oracle_impl import OracleDB
65
+
66
  return OracleDB
67
  elif db_type == "tidb":
68
  from ..kg.tidb_impl import TiDB
69
+
70
  return TiDB
71
  return None
72
 
73
+
74
  # Load environment variables
75
  try:
76
  load_dotenv(override=True)
 
908
  # Check which database types are used
909
  db_types = set()
910
  for storage_name, storage_instance in storage_instances:
911
+ db_type = get_db_type_from_storage_class(
912
+ storage_instance.__class__.__name__
913
+ )
914
  if db_type:
915
  db_types.add(db_type)
916
 
 
935
 
936
  # Inject database instances into storage classes
937
  for storage_name, storage_instance in storage_instances:
938
+ db_type = get_db_type_from_storage_class(
939
+ storage_instance.__class__.__name__
940
+ )
941
  if db_type:
942
  if db_type not in db_instances:
943
  error_msg = f"Database type '{db_type}' is required by {storage_name} but not initialized"
 
977
  db_names = {
978
  "postgres": "PostgreSQL",
979
  "oracle": "Oracle",
980
+ "tidb": "TiDB",
981
  }
982
  db_name = db_names.get(db_type, db_type)
983
  logger.info(f"Closed {db_name} database connection pool")
 
1300
  case ".pdf":
1301
  if not pm.is_installed("pypdf2"):
1302
  pm.install("pypdf2")
1303
+ from PyPDF2 import PdfReader # type: ignore
1304
  from io import BytesIO
1305
 
1306
  pdf_file = BytesIO(file)