yangdx
commited on
Commit
·
395a6bb
1
Parent(s):
a7c0d7b
Fix linting
Browse files
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(
|
|
|
|
|
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(
|
|
|
|
|
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
|
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)
|