fix(graph): 修复图谱文件上传的 bug,支持从MinIO URL加载知识图谱数据

添加从MinIO存储下载文件的功能,当检测到输入路径为URL时自动处理
下载完成后会创建临时文件并确保最终清理
This commit is contained in:
Wenjie Zhang 2025-12-15 11:01:00 +08:00
parent eb79ad4b8e
commit 3d6a38a73d

View File

@ -1,12 +1,15 @@
import json import json
import os import os
import tempfile
import traceback import traceback
import warnings import warnings
from urllib.parse import urlparse
from neo4j import GraphDatabase as GD from neo4j import GraphDatabase as GD
from src import config from src import config
from src.models import select_embedding_model from src.models import select_embedding_model
from src.storage.minio.client import get_minio_client, StorageError
from src.utils import logger from src.utils import logger
from src.utils.datetime_utils import utc_isoformat from src.utils.datetime_utils import utc_isoformat
@ -347,15 +350,59 @@ class GraphDatabase:
self.use_database(kgdb_name) # 切换到指定数据库 self.use_database(kgdb_name) # 切换到指定数据库
logger.info(f"Start adding entity to {kgdb_name} with {file_path}") logger.info(f"Start adding entity to {kgdb_name} with {file_path}")
def read_triples(file_path): # 检测 file_path 是否是 URL
with open(file_path, encoding="utf-8") as file: parsed_url = urlparse(file_path)
for line in file: temp_file_path = None
if line.strip():
yield json.loads(line.strip())
triples = list(read_triples(file_path)) try:
if parsed_url.scheme in ('http', 'https'): # 如果是 URL
logger.info(f"检测到 URL正在从 MinIO 下载文件: {file_path}")
await self.txt_add_vector_entity(triples, kgdb_name) # 从 URL 解析 bucket_name 和 object_name
# URL 格式: http://host:port/bucket_name/object_name
path_parts = parsed_url.path.lstrip('/').split('/', 1)
if len(path_parts) < 2:
raise ValueError(f"无法解析 MinIO URL: {file_path}")
bucket_name = path_parts[0]
object_name = path_parts[1]
# 从 MinIO 下载文件
minio_client = get_minio_client()
file_data = await minio_client.adownload_file(bucket_name, object_name)
# 创建临时文件保存下载的内容
with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False, encoding='utf-8') as temp_file:
temp_file.write(file_data.decode('utf-8'))
temp_file_path = temp_file.name
logger.info(f"文件已下载到临时路径: {temp_file_path}")
actual_file_path = temp_file_path
else:
# 本地文件路径
actual_file_path = file_path
def read_triples(file_path):
with open(file_path, encoding="utf-8") as file:
for line in file:
if line.strip():
yield json.loads(line.strip())
triples = list(read_triples(actual_file_path))
await self.txt_add_vector_entity(triples, kgdb_name)
except Exception as e:
logger.error(f"处理文件失败: {e}")
raise
finally:
# 清理临时文件
if temp_file_path and os.path.exists(temp_file_path):
try:
os.unlink(temp_file_path)
logger.info(f"已删除临时文件: {temp_file_path}")
except Exception as e:
logger.warning(f"删除临时文件失败: {e}")
self.status = "open" self.status = "open"
# 更新并保存图数据库信息 # 更新并保存图数据库信息