313 lines
10 KiB
Python
313 lines
10 KiB
Python
from appPublic.timeUtils import curDateString
|
||
from filemgr.filemgr import FileMgr
|
||
from rag.uapi_service import APIService
|
||
from appPublic.registerfunction import RegisterFunction
|
||
from appPublic.log import debug, error, info
|
||
from sqlor.dbpools import DBPools
|
||
import asyncio
|
||
import aiohttp
|
||
from langchain_core.documents import Document
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
import os
|
||
import re
|
||
import time
|
||
import uuid
|
||
from datetime import datetime
|
||
import traceback
|
||
from filetxt.loader import fileloader,File2Text
|
||
from ahserver.serverenv import get_serverenv
|
||
from typing import List, Dict, Any
|
||
from rag.service_opts import get_service_params, sor_get_service_params, sor_get_embedding_mode, get_embedding_mode
|
||
from rag.fileprocess import extract_images_from_file
|
||
from rag.rag_operations import RagOperations
|
||
import json
|
||
from rag.transaction_manager import TransactionContext
|
||
from dataclasses import dataclass
|
||
from enum import Enum
|
||
import base64
|
||
from pathlib import Path
|
||
|
||
class RagFileMgr(FileMgr):
|
||
def __init__(self, fiid):
|
||
super().__init__(fiid)
|
||
self.rag_ops = RagOperations()
|
||
|
||
async def get_folder_ownerid(self, sor):
|
||
fiid = self.fiid
|
||
recs = await sor.R('kdb', {'id': self.fiid})
|
||
if len(recs) > 0:
|
||
return recs[0].orgid
|
||
return None
|
||
|
||
async def get_organization_quota(self, sor, orgid):
|
||
sql = """select a.* from ragquota a, kdb b
|
||
where a.orgid = b.orgid
|
||
and b.id = ${id}$
|
||
and ${today}$ >= a.enabled_date
|
||
and ${today}$ < a.expired_date
|
||
"""
|
||
recs = await sor.sqlExe(sql, {
|
||
'id': self.fiid,
|
||
'today': curDateString()
|
||
})
|
||
if len(recs) > 0:
|
||
r = recs[0]
|
||
return r.quota, r.expired_date
|
||
return None, None
|
||
|
||
async def file_to_base64(self,path: str) -> str:
|
||
with open(path, "rb") as f:
|
||
return base64.b64encode(f.read()).decode("utf-8")
|
||
|
||
async def file_uploaded(self, request, ns, userid):
|
||
"""将文档插入 Milvus 并抽取三元组到 Neo4j"""
|
||
debug(f'Received ns: {ns=}')
|
||
env = request._run_ns
|
||
realpath = ns.get('realpath', '')
|
||
fiid = ns.get('fiid', '')
|
||
id = ns.get('id', '')
|
||
orgid = ns.get('ownerid', '')
|
||
db_type = ''
|
||
|
||
debug(
|
||
f'Inserting document: file_path={realpath}, userid={orgid}, db_type={db_type}, knowledge_base_id={fiid}, document_id={id}')
|
||
|
||
timings = {}
|
||
start_total = time.time()
|
||
result = {
|
||
"status": "error",
|
||
"userid": orgid,
|
||
"document_id": id,
|
||
"collection_name": "ragdb",
|
||
"timings": timings,
|
||
"message": "",
|
||
"status_code": 400
|
||
}
|
||
|
||
# 初始化回滚上下文
|
||
rollback_context = {
|
||
"request": request,
|
||
"userid": userid,
|
||
"service_params": None # 在 try 块中设置
|
||
}
|
||
|
||
async with TransactionContext(f"file_upload_{id}") as transaction_mgr:
|
||
# 将 rollback_context 绑定到 TransactionContext
|
||
transaction_mgr.transaction_context = rollback_context
|
||
try:
|
||
# 验证必填字段
|
||
if not orgid or not fiid or not id:
|
||
raise ValueError("orgid、fiid 和 id 不能为空")
|
||
if len(orgid) > 32 or len(fiid) > 255:
|
||
raise ValueError("orgid 或 fiid 的长度超出限制")
|
||
if not os.path.exists(realpath):
|
||
raise ValueError(f"文件 {realpath} 不存在")
|
||
|
||
# 获取服务参数
|
||
service_params = await get_service_params(orgid)
|
||
debug(f"服务参数是：{service_params}")
|
||
if not service_params:
|
||
raise ValueError("无法获取服务参数")
|
||
rollback_context["service_params"] = service_params
|
||
|
||
#获取嵌入模式
|
||
embedding_mode = await get_embedding_mode(orgid)
|
||
debug(f"检测到 embedding_mode = {embedding_mode}（0=文本, 1=多模态）")
|
||
|
||
# 加载和分片文档
|
||
chunks = await self.rag_ops.load_and_chunk_document(
|
||
realpath, timings, transaction_mgr=transaction_mgr
|
||
)
|
||
|
||
text_embeddings = None
|
||
multi_results = None
|
||
image_paths = []
|
||
|
||
if embedding_mode == 1:
|
||
inputs = []
|
||
# 文本
|
||
for chunk in chunks:
|
||
inputs.append({"type": "text", "content": chunk.page_content})
|
||
|
||
debug("开始多模态图像抽取与嵌入")
|
||
image_paths = extract_images_from_file(realpath)
|
||
debug(f"从文档中抽取 {len(image_paths)} 张图像")
|
||
|
||
if image_paths:
|
||
for img_path in image_paths:
|
||
try:
|
||
# 1. 自动识别真实格式
|
||
ext = Path(img_path).suffix.lower()
|
||
if ext not in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}:
|
||
ext = ".jpg"
|
||
|
||
mime_map = {
|
||
".png": "image/png",
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".webp": "image/webp",
|
||
".bmp": "image/bmp"
|
||
}
|
||
mime_type = mime_map.get(ext, "image/jpeg")
|
||
|
||
# # 2. 智能压缩（>1MB 才压缩，节省 70% 流量）
|
||
# img = Image.open(img_path).convert("RGB")
|
||
# if os.path.getsize(img_path) > 1024 * 1024: # >1MB
|
||
# buffer = BytesIO()
|
||
# img.save(buffer, format="JPEG", quality=85, optimize=True)
|
||
# b64 = base64.b64encode(buffer.getvalue()).decode()
|
||
# data_uri = f"data:image/jpeg;base64,{b64}"
|
||
# else:
|
||
b64 = await self.file_to_base64(img_path)
|
||
data_uri = f"data:{mime_type};base64,{b64}"
|
||
|
||
inputs.append({
|
||
"type": "image",
|
||
"data": data_uri
|
||
})
|
||
debug(f"已添加图像（{mime_type}, {len(b64) / 1024:.1f}KB）: {Path(img_path).name}")
|
||
|
||
except Exception as e:
|
||
debug(f"图像处理失败，跳过: {img_path} → {e}")
|
||
# 即使失败也加个占位，防止顺序错乱
|
||
inputs.append({
|
||
"type": "image",
|
||
"data": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
|
||
})
|
||
|
||
debug(f"混排输入总数: {len(inputs)}（文本 {len(chunks)} + 图像 {len(image_paths)}）")
|
||
|
||
multi_results = await self.rag_ops.generate_multi_embeddings(
|
||
request=request,
|
||
inputs=inputs,
|
||
service_params=service_params,
|
||
userid=userid,
|
||
timings=timings,
|
||
transaction_mgr=transaction_mgr
|
||
)
|
||
debug(f"多模态嵌入成功，返回 {len(multi_results)} 条结果")
|
||
else:
|
||
# 生成嵌入向量
|
||
debug("【纯文本模式】使用 BGE 嵌入")
|
||
text_embeddings = await self.rag_ops.generate_embeddings(
|
||
request, chunks, service_params, userid, timings, transaction_mgr=transaction_mgr
|
||
)
|
||
debug(f"BGE 嵌入完成: {len(text_embeddings)} 条")
|
||
|
||
inserted = await self.rag_ops.insert_all_vectors(
|
||
request=request,
|
||
text_chunks=chunks,
|
||
realpath=realpath,
|
||
orgid=orgid,
|
||
fiid=fiid,
|
||
document_id=id,
|
||
service_params=service_params,
|
||
userid=userid,
|
||
db_type=db_type,
|
||
timings=timings,
|
||
img_paths=image_paths,
|
||
text_embeddings=text_embeddings,
|
||
multi_results=multi_results,
|
||
transaction_mgr=transaction_mgr
|
||
)
|
||
debug(f"统一插入: 文本 {inserted['text']}, 图像 {inserted['image']}, 人脸 {inserted['face']}")
|
||
|
||
# 抽取三元组
|
||
triples = await self.rag_ops.extract_triples(
|
||
request, chunks, service_params, userid, timings, transaction_mgr=transaction_mgr
|
||
)
|
||
|
||
# 插入 Neo4j
|
||
await self.rag_ops.insert_to_graph_db(
|
||
request, triples, id, fiid, orgid, service_params, userid, timings, transaction_mgr=transaction_mgr
|
||
)
|
||
|
||
timings["total"] = time.time() - start_total
|
||
result.update({
|
||
"status": "success",
|
||
"unique_triples": triples,
|
||
"message": f"文件 {realpath} 成功嵌入并处理三元组",
|
||
"status_code": 200
|
||
})
|
||
debug(f"总耗时: {timings['total']:.2f} 秒")
|
||
|
||
except Exception as e:
|
||
error(f"插入文档失败: {str(e)}, 堆栈: {traceback.format_exc()}")
|
||
timings["total"] = time.time() - start_total
|
||
result.update({
|
||
"message": f"插入文档失败: {str(e)}",
|
||
"timings": timings
|
||
})
|
||
raise ValueError(str(e)) from e
|
||
debug(f"最终结果是：{result}")
|
||
return result
|
||
|
||
async def file_deleted(self, request, recs, userid):
|
||
"""删除用户指定文件数据，包括 Milvus 和 Neo4j 中的记录"""
|
||
if not isinstance(recs, list):
|
||
recs = [recs]
|
||
results = []
|
||
total_nodes_deleted = 0
|
||
total_rels_deleted = 0
|
||
|
||
for rec in recs:
|
||
id = rec.get('id', '')
|
||
realpath = rec.get('realpath', '')
|
||
fiid = rec.get('fiid', '')
|
||
orgid = rec.get('ownerid', '')
|
||
db_type = ''
|
||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||
|
||
try:
|
||
required_fields = ['id', 'realpath', 'fiid', 'ownerid']
|
||
missing_fields = [field for field in required_fields if not rec.get(field, '')]
|
||
if missing_fields:
|
||
raise ValueError(f"缺少必填字段: {', '.join(missing_fields)}")
|
||
|
||
service_params = await get_service_params(orgid)
|
||
if not service_params:
|
||
raise ValueError("无法获取服务参数")
|
||
|
||
# 调用 Milvus 删除
|
||
await self.rag_ops.delete_from_vector_db(request, orgid, realpath, fiid, id, service_params, userid, db_type)
|
||
|
||
# 调用 Neo4j 删除
|
||
neo4j_deleted_nodes = 0
|
||
neo4j_deleted_rels = 0
|
||
try:
|
||
nodes_deleted, rels_deleted = await self.rag_ops.delete_from_graph_db(request, id, service_params, userid)
|
||
neo4j_deleted_nodes += nodes_deleted
|
||
neo4j_deleted_rels += rels_deleted
|
||
total_nodes_deleted += nodes_deleted
|
||
total_rels_deleted += rels_deleted
|
||
except Exception as e:
|
||
error(f"删除 document_id={id} 的 Neo4j 数据失败: {str(e)}")
|
||
|
||
results.append({
|
||
"status": "success",
|
||
"collection_name": collection_name,
|
||
"document_id": id,
|
||
"message": f"成功删除文件 {realpath} 的 Milvus 记录，{neo4j_deleted_nodes} 个 Neo4j 节点，{neo4j_deleted_rels} 个 Neo4j 关系",
|
||
"status_code": 200
|
||
})
|
||
|
||
except Exception as e:
|
||
error(f"删除文档 {realpath} 失败: {str(e)}, 堆栈: {traceback.format_exc()}")
|
||
results.append({
|
||
"status": "error",
|
||
"collection_name": collection_name,
|
||
"document_id": id,
|
||
"message": f"删除文档 {realpath} 失败: {str(e)}",
|
||
"status_code": 400
|
||
})
|
||
|
||
return {
|
||
"status": "success" if all(r["status"] == "success" for r in results) else "partial",
|
||
"results": results,
|
||
"total_nodes_deleted": total_nodes_deleted,
|
||
"total_rels_deleted": total_rels_deleted,
|
||
"message": f"处理 {len(recs)} 个文件，成功删除 {sum(1 for r in results if r['status'] == 'success')} 个",
|
||
"status_code": 200 if all(r["status"] == "success" for r in results) else 207
|
||
}
|
||
|