KbmService.py 65 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705
  1. # 导入 backend/__init__.py 中的所有内容
  2. from .. import *
  3. # # 模型
  4. from backend.models import Knowledgebase, DocumentKbm, File2document, File, Task, TaskSublist, KbmDocumentType
  5. from django.db.models import Count, Case, When, IntegerField, Q, Max
  6. from django.core.paginator import Paginator, PageNotAnInteger, EmptyPage
  7. from django.shortcuts import get_object_or_404
  8. from django.db import transaction
  9. from django.core.exceptions import ObjectDoesNotExist
  10. from backend.Service.MinioService import MinioService
  11. from tabulate import tabulate
  12. from scipy.spatial.distance import cosine
  13. from minio import Minio
  14. from django.conf import settings
  15. import pytesseract
  16. from DCbackend.utils.common import success, fail
  17. from DCbackend.settings import MILVUS_HOST, MILVUS_PORT, VECTOR_DIMENSION, IPINFO, MILVUS_USER, MILVUS_PASSWORD,MINIO_ENDPOINT,MINIO_ACCESS_KEY,MINIO_SECRET_KEY,MINIO_SECURE
  18. # 导入pymilvus库中的连接、集合、字段模式、集合模式、数据类型和实用工具
  19. from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
  20. from bs4 import BeautifulSoup
  21. # 从service_utils模块中导入rabbitmq和pdf_utils
  22. from .service_utils import rabbitmq as rabbitmq_Process , pdf_utils,ocr_utils
  23. import chardet
  24. import threading # 基于线程的并行处理,用于多线程编程
  25. os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
  26. class DocumentQueue:
  27. QUEUE_KEY = "document_process_queue"
  28. minio_client = Minio(
  29. MINIO_ENDPOINT,
  30. access_key=MINIO_ACCESS_KEY,
  31. secret_key=MINIO_SECRET_KEY,
  32. secure=MINIO_SECURE # 根据您的 Minio 配置选择 True 或 False
  33. )
  34. if os.name == 'nt': # Windows
  35. pytesseract.pytesseract.tesseract_cmd = r'D:\Program Files\OCR\tesseract.exe'
  36. else: # macOS 或 Linux
  37. pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract'
  38. class KbmService:
  39. NLP_ZH = spacy.load('zh_core_web_sm')
  40. bert_model = None
  41. bert_tokenizer = None
  42. # 大模型地址
  43. API_URL = f"http://{IPINFO}:11434/api/embeddings"
  44. semaphore = threading.Semaphore(4)
  45. # 通用rabbitmq 放入队列
  46. def send_to_rabbitmq(queue_name, message):
  47. """
  48. 将消息发送到指定的RabbitMQ队列
  49. Args:
  50. queue_name (str): 目标RabbitMQ队列的名称
  51. message (dict): 需要发送的消息
  52. Returns:
  53. bool: 消息发送成功返回True,失败返回False
  54. """
  55. return rabbitmq_Process.send_to_rabbitmq(queue_name, message)
  56. @staticmethod
  57. def selectBucketInfo(request):
  58. """
  59. 从请求中获取知识库信息,包括文档数量统计
  60. Args:
  61. request (HttpRequest): 请求对象,包含查询条件
  62. Returns:
  63. dict: 成功时返回包含所有知识库信息和文档计数的字典,失败时返回包含错误信息的字典
  64. """
  65. #user_id = request.POST.get("user_id")
  66. knowledgebases = Knowledgebase.objects.filter().exclude(status=4).order_by('-create_time').values('id', 'create_time', 'name', 'doc_num', 'description')
  67. result = [] # 初始化结果列表
  68. for kb in knowledgebases:
  69. # 对每个知识库,统计不同类型文档的数量
  70. counts = DocumentKbm.objects.filter(kb_id=kb['id']).aggregate(
  71. word_count=Count(Case(When(type__in=['doc', 'docx'], then=1), output_field=IntegerField())),
  72. pdf_count=Count(Case(When(type='pdf', then=1), output_field=IntegerField())),
  73. excel_count=Count(Case(When(type__in=['xls', 'xlsx'], then=1), output_field=IntegerField()))
  74. )
  75. # 创建知识库数据的副本,以避免修改原始数据
  76. kb_data = kb.copy()
  77. # 将统计的计数更新到知识库数据副本中
  78. kb_data.update(counts)
  79. # 将更新后的知识库数据添加到结果列表中
  80. result.append(kb_data)
  81. # 返回包含所有知识库信息和文档计数的成功响应
  82. return success(result)
  83. @staticmethod
  84. def getFileInfo(request):
  85. """
  86. 获取文件信息
  87. 从请求的POST数据中获取bucket_id、page、pageSize、object_name、run、type和doc_type_id,用于查询和分页文档信息。
  88. 查询条件包括知识库ID、文档名称、运行标识、文档类型和文档类型ID。结果包括文档信息和分页信息。
  89. Args:
  90. request (HttpRequest): 请求对象,包含查询条件。
  91. Returns:
  92. dict: 成功时返回包含文档信息和分页信息的字典,失败时返回包含错误信息的字典。
  93. """
  94. try:
  95. bucket_id = request.POST.get("bucket_id")
  96. page = request.POST.get("page", 1)
  97. per_page = request.POST.get("pageSize", 10)
  98. object_name = request.POST.get("object_name", "")
  99. run = request.POST.get("run", "")
  100. type = request.POST.get("type", "")
  101. doc_type_id = request.POST.get("doc_type_id")
  102. if not bucket_id:
  103. return fail("bucket_id为空")
  104. # 确保 page 和 per_page 是整数
  105. page = int(page)
  106. per_page = int(per_page)
  107. # # 查询文档并排序
  108. # documents = DocumentKbm.objects.filter(
  109. # Q(kb_id=bucket_id) &
  110. # Q(name__icontains=object_name)&
  111. # Q(run__icontains=run)&
  112. # Q(type__icontains=type)&
  113. # Q(doc_type_id=doc_type_id)&
  114. # ~Q(status=4)
  115. # ).order_by('-create_time')
  116. # 构建基本查询
  117. query = Q(kb_id=bucket_id) & Q(name__icontains=object_name) & Q(run__icontains=run) & Q(type__icontains=type) & ~Q(status=4)
  118. # 如果 doc_type_id 有值,则添加到查询条件
  119. if doc_type_id:
  120. query &= Q(doc_type_id=doc_type_id)
  121. # 查询文档并排序
  122. documents = DocumentKbm.objects.filter(query).order_by('-create_time')
  123. # 创建分页器
  124. paginator = Paginator(documents, per_page)
  125. try:
  126. # 获取指定页的结果
  127. documents_page = paginator.page(page)
  128. except PageNotAnInteger:
  129. # 如果页码不是整数,返回第一页
  130. documents_page = paginator.page(1)
  131. except EmptyPage:
  132. # 如果页码超出范围,返回最后一页
  133. documents_page = paginator.page(paginator.num_pages)
  134. # 将查询结果转换为列表
  135. result = list(documents_page.object_list.values())
  136. # 循环整体会导致mysql多连反应,(后期应该改成数据载加到内存环境后一步塞入数据库中)
  137. # 发现人 qman
  138. for info in result:
  139. document_id = info['id']
  140. max_page = TaskSublist.objects.filter(doc_id=document_id).aggregate(Max('page_number'))['page_number__max']
  141. info['max_page'] = max_page if max_page is not None else 0
  142. pagination_info = {
  143. 'total_count': paginator.count,
  144. 'total_pages': paginator.num_pages,
  145. 'total_size': per_page,
  146. 'current_page': documents_page.number,
  147. 'has_next': documents_page.has_next(),
  148. 'has_previous': documents_page.has_previous()
  149. }
  150. data = {
  151. 'pagination': pagination_info,
  152. 'documents': result
  153. }
  154. return success(data)
  155. except Exception as e:
  156. return fail("获取信息失败")
  157. @staticmethod
  158. @transaction.atomic
  159. def updateName(request):
  160. """
  161. 更新文档的名称
  162. 从请求的POST数据中获取新名称和文档ID,更新对应的DocumentKbm对象和关联的File对象的名称。
  163. 如果新名称不包含扩展名,则使用原始文件的扩展名。
  164. Args:
  165. request (HttpRequest): 请求对象,包含新名称和文档ID。
  166. Returns:
  167. dict: 成功时返回包含成功信息的字典,失败时返回包含错误信息的字典。
  168. """
  169. try:
  170. new_name = request.POST.get("new_name")
  171. document_id = request.POST.get("document_id")
  172. if not new_name or not document_id:
  173. return fail("新名称和文件ID不能为空")
  174. # 获取 DocumentKbm 实例并更新
  175. document = get_object_or_404(DocumentKbm, id=document_id)
  176. location = document.location
  177. # 获取原始文件的扩展名
  178. _, original_extension = os.path.splitext(location)
  179. # 检查新名称是否包含扩展名,如果没有则添加原始扩展名
  180. _, new_extension = os.path.splitext(new_name)
  181. if not new_extension:
  182. new_name = f"{new_name}{original_extension}"
  183. document.name = new_name
  184. document.save()
  185. # 获取关联的 File2document 和 File
  186. file2doc = File2document.objects.filter(document_id=document_id).first()
  187. if file2doc:
  188. file = get_object_or_404(File, id=file2doc.file_id)
  189. file.name = new_name
  190. file.save()
  191. else:
  192. # 记录一个警告,因为没有找到关联的 File
  193. logger.error(f"Warning: No associated File found for DocumentKbm with id {document_id}")
  194. return success("文件名更新成功")
  195. except ObjectDoesNotExist:
  196. return fail("指定的文件或关联文件不存在")
  197. except Exception as e:
  198. return fail(f"更新文件名失败: {str(e)}")
  199. @staticmethod
  200. @transaction.atomic
  201. def deleteDocument(request):
  202. """
  203. 删除文档
  204. 从请求的POST数据中获取文档ID,更新文档和关联文件的状态为已删除。然后,更新知识库的文档数量。
  205. 最后,尝试清理Milvus中的相关数据,并断开Milvus连接。
  206. Args:
  207. request (HttpRequest): 请求对象,包含文档ID。
  208. Returns:
  209. dict: 成功时返回包含成功信息的字典,失败时返回包含错误信息的字典。
  210. """
  211. document_id = request.POST.get("document_id")
  212. document = get_object_or_404(DocumentKbm, id=document_id)
  213. document.status = 4
  214. document.save()
  215. file2doc = File2document.objects.filter(document_id=document_id).first()
  216. if file2doc:
  217. file = get_object_or_404(File, id=file2doc.file_id)
  218. file.status = 4
  219. file.save()
  220. kb_id = document.kb_id
  221. new_count = DocumentKbm.objects.filter(kb_id=kb_id).exclude(status=4).count()
  222. Knowledgebase.objects.filter(id=kb_id).update(doc_num=new_count)
  223. try:
  224. # 需要 拆列方法 石峰
  225. # 1 、 删除 minio
  226. # 2、 删除 milvus
  227. # 清理milvus
  228. # 连接到 Milvus
  229. connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT ,user=MILVUS_USER,password=MILVUS_PASSWORD)
  230. kmb = Knowledgebase.objects.filter(id=kb_id).first()
  231. tasks = TaskSublist.objects.filter(doc_id=document.id)
  232. KbmService.clearPreviousData(document_id,document.name)
  233. # collection = Collection(kmb.location)
  234. # tasks = TaskSublist.objects.filter(doc_id=document.id)
  235. # for task in tasks:
  236. # logger.info(f'删除{task.milvus_id}')
  237. # expr = f'id in [{task.milvus_id}]'
  238. # collection.delete(expr)
  239. # minio_client.remove_object(kmb.location, tasks.name)
  240. # files = File.objects.filter(id=file2doc.file_id).first()
  241. # minio_client.remove_object(kmb.location, files.location)
  242. return success("删除成功")
  243. except Exception as e:
  244. return fail(f"删除milvus集合时发生错误: {str(e)}")
  245. finally:
  246. # 断开 Milvus 连接
  247. connections.disconnect("default")
  248. @staticmethod
  249. def getUrl(request):
  250. """
  251. 根据请求获取文档的URL。
  252. 从请求的POST数据中获取文档ID,通过文档ID获取对应的DocumentKbm对象和关联的Knowledgebase对象。
  253. 然后,使用MinioService.geturl方法获取文档的URL。
  254. Args:
  255. request (HttpRequest): 请求对象,包含文档ID。
  256. Returns:
  257. dict: 成功时返回包含文档URL的字典,失败时返回包含错误信息的字典。
  258. """
  259. try:
  260. document_id = request.POST.get("document_id")
  261. if not document_id:
  262. return fail("文档ID不能为空")
  263. # 获取 DocumentKbm 对象
  264. document = get_object_or_404(DocumentKbm, id=document_id)
  265. object_name = document.location
  266. # 获取对应的 Knowledgebase 对象
  267. knowledgebase = get_object_or_404(Knowledgebase, id=document.kb_id)
  268. bucket_name = knowledgebase.location
  269. return MinioService.geturl(object_name, bucket_name)
  270. except ObjectDoesNotExist:
  271. return fail("指定的文档或知识库不存在")
  272. except Exception as e:
  273. return fail(f"获取URL失败: {str(e)}")
  274. @staticmethod
  275. def get_embedding_excel(text, target_dim=768):
  276. """
  277. 根据输入文本获取其对应的嵌入向量。
  278. Args:
  279. text (str): 输入文本。
  280. target_dim (int, optional): 目标维度。默认为768。
  281. Returns:
  282. list: 嵌入向量的列表表示。
  283. """
  284. try:
  285. if not text or not text.strip():
  286. logger.warning("Empty text provided for embedding. Returning zero vector.")
  287. return np.zeros(target_dim).tolist()
  288. # 确保文本被正确编码
  289. encoded_text = text.encode('utf-8').decode('utf-8')
  290. payload = {
  291. "model": "nomic-embed-text:latest",
  292. "prompt": encoded_text
  293. }
  294. headers = {"Content-Type": "application/json"}
  295. #
  296. response = requests.post(KbmService.API_URL, json=payload, headers=headers)
  297. logger.info(f"response::::{response}")
  298. response.raise_for_status()
  299. embedding_data = response.json()
  300. if 'embedding' not in embedding_data:
  301. raise ValueError(f"API 响应中没有找到嵌入向量. 响应内容: {embedding_data}")
  302. embedding = embedding_data['embedding']
  303. original_embedding = np.array(embedding)
  304. if len(original_embedding) == target_dim:
  305. return original_embedding.tolist()
  306. # 如果原始维度不等于目标维度,进行插值
  307. original_indices = np.arange(len(original_embedding))
  308. new_indices = np.linspace(0, len(original_embedding) - 1, target_dim)
  309. f = interpolate.interp1d(original_indices, original_embedding)
  310. extended_embedding = f(new_indices)
  311. return extended_embedding.tolist()
  312. except requests.exceptions.RequestException as e:
  313. logger.error(f"API 请求错误: {str(e)}")
  314. raise
  315. except ValueError as e:
  316. logger.error(f"值错误: {str(e)}")
  317. raise
  318. except Exception as e:
  319. logger.error(f"获取文本嵌入时发生意外错误: {str(e)}")
  320. raise
  321. #新rabbitmq队列
  322. @staticmethod
  323. def analysis(request):
  324. """
  325. 分析请求并处理RabbitMQ队列中的消息。
  326. Args:
  327. request (object): 需要分析的请求对象。
  328. Returns:
  329. None
  330. """
  331. return rabbitmq_Process.analysis(request, KbmService)
  332. @staticmethod
  333. def check_and_process_queue():
  334. """
  335. 检查并处理RabbitMQ队列中的消息,启用多线程 。
  336. Returns:
  337. None
  338. """
  339. KbmService.should_stop = False
  340. while not KbmService.should_stop:
  341. try:
  342. if KbmService.queue_has_messages():
  343. KbmService.process_queue()
  344. else:
  345. # logger.info("队列为空,等待下一次检查...")
  346. time.sleep(60) # 等待60秒后再次检查
  347. except Exception as e:
  348. logger.error(f"检查队列时发生错误: {str(e)}")
  349. time.sleep(60) # 发生错误时,等待60秒后重试
  350. @staticmethod
  351. def queue_has_messages():
  352. """
  353. 检查RabbitMQ队列中是否有消息。
  354. Returns:
  355. bool: 如果队列中有消息返回True,否则返回False。
  356. """
  357. try:
  358. connection = KbmService.create_connection()
  359. channel = connection.channel()
  360. queue = channel.queue_declare(queue=settings.RABBITMQ_QUEUE_NAME, passive=True)
  361. message_count = queue.method.message_count
  362. connection.close()
  363. return message_count > 0
  364. except Exception as e:
  365. logger.error(f"检查队列消息数量时发生错误: {str(e)}")
  366. return False
  367. @staticmethod
  368. def create_connection():
  369. """
  370. 创建到RabbitMQ服务器的连接。
  371. Returns:
  372. pika.BlockingConnection: 连接对象。
  373. """
  374. return pika.BlockingConnection(pika.ConnectionParameters(
  375. host=settings.RABBITMQ_HOST,
  376. port=settings.RABBITMQ_PORT,
  377. credentials=pika.PlainCredentials(
  378. settings.RABBITMQ_USER,
  379. settings.RABBITMQ_PASSWORD
  380. )
  381. ))
  382. @staticmethod
  383. def process_queue():
  384. logger.info("队列中有消息,开始处理...")
  385. KbmService.connection = KbmService.create_connection()
  386. KbmService.channel = KbmService.connection.channel()
  387. KbmService.channel.queue_declare(queue=settings.RABBITMQ_QUEUE_NAME, durable=True)
  388. KbmService.channel.basic_qos(prefetch_count=4)
  389. KbmService.channel.basic_consume(queue=settings.RABBITMQ_QUEUE_NAME, on_message_callback=KbmService.callback)
  390. try:
  391. KbmService.channel.start_consuming()
  392. except KeyboardInterrupt:
  393. KbmService.should_stop = True
  394. finally:
  395. KbmService.close_connection()
  396. @staticmethod
  397. def callback(ch, method, properties, body):
  398. """
  399. 处理RabbitMQ队列中的消息。
  400. Args:
  401. ch (pika.channel.Channel): RabbitMQ频道对象。
  402. method (pika.spec.Basic.Deliver): 消息投递信息。
  403. properties (pika.spec.BasicProperties): 消息属性。
  404. body (bytes): 消息体。
  405. Raises:
  406. Exception: 处理消息时发生的任何异常。
  407. """
  408. with KbmService.semaphore:
  409. try:
  410. job = json.loads(body)
  411. document_id = job['document_id']
  412. start_page = job['start_page']
  413. end_page = job['end_page']
  414. max_tokens = job['max_tokens']
  415. logger.info(f"开始执行解析文档 {document_id}")
  416. KbmService.async_analysis(document_id, start_page, end_page, max_tokens)
  417. ch.basic_ack(delivery_tag=method.delivery_tag)
  418. except Exception as e:
  419. logger.error(f"处理队列消息时发生错误: {str(e)}")
  420. ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
  421. # 检查是否还有更多消息
  422. if not KbmService.queue_has_messages():
  423. logger.info("队列处理完毕,停止消费...")
  424. ch.stop_consuming()
  425. @staticmethod
  426. def close_connection():
  427. if KbmService.channel:
  428. try:
  429. KbmService.channel.close()
  430. except Exception:
  431. pass
  432. if KbmService.connection:
  433. try:
  434. KbmService.connection.close()
  435. except Exception:
  436. pass
  437. KbmService.channel = None
  438. KbmService.connection = None
  439. @staticmethod
  440. def stop_service():
  441. KbmService.should_stop = True
  442. if KbmService.channel:
  443. KbmService.channel.stop_consuming()
  444. KbmService.close_connection()
  445. @staticmethod
  446. def get_embedding_excel(text, target_dim=768):
  447. """
  448. 获取文本的嵌入向量
  449. 从API获取文本的嵌入向量,并根据需要调整到目标维度。
  450. Args:
  451. text (str): 需要获取嵌入向量的文本。
  452. target_dim (int, optional): 目标维度。默认为768。
  453. Returns:
  454. list: 文本的嵌入向量列表,长度为target_dim。
  455. Raises:
  456. Exception: 处理文本嵌入时发生的任何异常。
  457. """
  458. try:
  459. if not text or not text.strip():
  460. logger.warning("Empty text provided for embedding. Returning zero vector.")
  461. return np.zeros(target_dim).tolist()
  462. # 确保文本被正确编码
  463. encoded_text = text.encode('utf-8').decode('utf-8')
  464. payload = {
  465. "model": "nomic-embed-text:latest",
  466. "prompt": encoded_text
  467. }
  468. headers = {"Content-Type": "application/json"}
  469. response = requests.post(KbmService.API_URL, json=payload, headers=headers)
  470. logger.info(f"response::::{response}")
  471. response.raise_for_status()
  472. embedding_data = response.json()
  473. if 'embedding' not in embedding_data:
  474. raise ValueError(f"API 响应中没有找到嵌入向量. 响应内容: {embedding_data}")
  475. embedding = embedding_data['embedding']
  476. original_embedding = np.array(embedding)
  477. if len(original_embedding) == target_dim:
  478. return original_embedding.tolist()
  479. # 如果原始维度不等于目标维度,进行插值
  480. original_indices = np.arange(len(original_embedding))
  481. new_indices = np.linspace(0, len(original_embedding) - 1, target_dim)
  482. f = interpolate.interp1d(original_indices, original_embedding)
  483. extended_embedding = f(new_indices)
  484. return extended_embedding.tolist()
  485. except requests.exceptions.RequestException as e:
  486. logger.error(f"API 请求错误: {str(e)}")
  487. raise
  488. except ValueError as e:
  489. logger.error(f"值错误: {str(e)}")
  490. raise
  491. except Exception as e:
  492. logger.error(f"获取文本嵌入时发生意外错误: {str(e)}")
  493. raise
  494. @classmethod
  495. def get_embedding_pdf(cls, text, target_dim=768, max_retries=3, backoff_factor=0.3):
  496. """
  497. 获取文本的嵌入向量,并填充或截断到目标维度,包含重试机制
  498. Args:
  499. text (str): 需要嵌入的文本。
  500. target_dim (int, optional): 目标维度。默认为 768。
  501. max_retries (int, optional): 最大重试次数。默认为 3。
  502. backoff_factor (float, optional): 每次重试的退避因子。默认为 0.3。
  503. Returns:
  504. np.ndarray: 填充或截断后的嵌入向量。
  505. """
  506. # 清楚 一些 垃圾 字符
  507. text = KbmService.post_process_text(text)
  508. for attempt in range(max_retries):
  509. try:
  510. payload = {
  511. "model": "nomic-embed-text:latest",
  512. "prompt": text
  513. }
  514. headers = {
  515. "Content-Type": "application/json"
  516. }
  517. response = requests.post(cls.API_URL, json=payload, headers=headers, timeout=30)
  518. sleep(0.3)
  519. response.raise_for_status()
  520. result = response.json()
  521. embedding = result.get('embedding')
  522. if embedding is None:
  523. raise ValueError("API 响应中没有找到嵌入向量")
  524. embedding_array = np.array(embedding)
  525. current_dim = embedding_array.shape[0]
  526. if current_dim < target_dim:
  527. padded_embedding = np.pad(embedding_array, (0, target_dim - current_dim), 'constant')
  528. logger.info(f"向量已从 {current_dim} 维填充到 {target_dim} 维")
  529. return padded_embedding
  530. elif current_dim > target_dim:
  531. truncated_embedding = embedding_array[:target_dim]
  532. logger.info(f"向量已从 {current_dim} 维截断到 {target_dim} 维")
  533. return truncated_embedding
  534. else:
  535. return embedding_array
  536. except RequestException as e:
  537. logger.error(f"API 请求错误 (尝试 {attempt + 1}/{max_retries}): {e}")
  538. if attempt == max_retries - 1:
  539. raise
  540. time.sleep(backoff_factor * (2 ** attempt))
  541. except ValueError as e:
  542. logger.error(f"解析响应错误: {e}")
  543. raise
  544. except Exception as e:
  545. logger.error(f"获取文本嵌入时发生未知错误: {e}")
  546. raise
  547. raise Exception("达到最大重试次数,无法获取嵌入")
  548. @staticmethod
  549. def split_text_by_semantic(sentences, max_tokens, bucket_name, similarity_threshold=0.5, batch_size=1000):
  550. logger.info("开始分割文本并保存到向量数据库")
  551. chunks = []
  552. if len(sentences)<= 0:
  553. return chunks
  554. try:
  555. #object1 object1为后续可能添加的字段 因为无法直接修改名称 备用
  556. source = "知识库"
  557. object1 ="some_object1"
  558. object2 ="some_object2"
  559. # 连接到Milvus
  560. connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT,user=MILVUS_USER,password=MILVUS_PASSWORD)
  561. collection_name = f"{bucket_name}"
  562. collection = KbmService._get_or_create_collection(collection_name)
  563. current_chunk = sentences[0]
  564. current_embedding = KbmService.get_embedding_pdf(current_chunk, target_dim=VECTOR_DIMENSION)
  565. batch_data = []
  566. for sentence in sentences[1:]:
  567. if len(sentence) > 10:
  568. sentence_embedding = KbmService.get_embedding_pdf(sentence, target_dim=VECTOR_DIMENSION)
  569. similarity = 1 - cosine(current_embedding, sentence_embedding)
  570. if len(current_chunk) + len(sentence) <= max_tokens and similarity >= similarity_threshold:
  571. current_chunk += sentence
  572. current_embedding = (current_embedding + sentence_embedding) / 2
  573. else:
  574. batch_data.append((current_chunk, current_embedding))
  575. if len(batch_data) >= batch_size:
  576. ids = KbmService._insert_batch(collection, batch_data,source,object1,object2)
  577. sleep(1)
  578. logger.info("减少milvus压力睡眠1秒")
  579. if ids is not None:
  580. chunks.extend([{'content': chunk, 'milvus_id': id} for (chunk, _), id in zip(batch_data, ids)])
  581. else:
  582. logger.error("向 Milvus 插入批量数据失败,这批数据将被跳过")
  583. batch_data = []
  584. current_chunk = sentence
  585. current_embedding = sentence_embedding
  586. # 处理最后一个chunk和剩余的batch数据
  587. if current_chunk:
  588. batch_data.append((current_chunk, current_embedding))
  589. if batch_data:
  590. ids = KbmService._insert_batch(collection, batch_data,source,object1,object2)
  591. sleep(1)
  592. logger.info("减少milvus压力睡眠1秒")
  593. if ids is not None:
  594. chunks.extend([{'content': chunk, 'milvus_id': id} for (chunk, _), id in zip(batch_data, ids)])
  595. else:
  596. logger.error("向 Milvus 插入批量数据失败,这批数据将被跳过")
  597. KbmService._create_index_and_load(collection)
  598. logger.info(f"成功将{len(chunks)}个文本块分割并保存到Milvus")
  599. return chunks
  600. except Exception as e:
  601. logger.error(f"处理文本时发生错误: {str(e)}")
  602. raise
  603. finally:
  604. connections.disconnect("default")
  605. @staticmethod
  606. def _get_or_create_collection(collection_name):
  607. if not utility.has_collection(collection_name):
  608. fields = [
  609. FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
  610. FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=65000),
  611. FieldSchema(name="object1", dtype=DataType.VARCHAR, max_length=65000),
  612. FieldSchema(name="object2", dtype=DataType.VARCHAR, max_length=65000),
  613. FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65000),
  614. FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=VECTOR_DIMENSION)
  615. ]
  616. schema = CollectionSchema(fields, "Semantic text chunks collection")
  617. return Collection(name=collection_name, schema=schema)
  618. return Collection(name=collection_name)
  619. @staticmethod
  620. def _split_sentences(text):
  621. sentences = re.split('([。!?])', text)
  622. sentences = [''.join(i) for i in zip(sentences[0::2], sentences[1::2] + [''])]
  623. return [s.strip() for s in sentences if s.strip()]
  624. @staticmethod
  625. def _insert_batch(collection, batch_data, source, object1, object2):
  626. try:
  627. entities = [
  628. [source] * len(batch_data), # source
  629. [object1]* len(batch_data), # object1
  630. [object2]* len(batch_data), # object2
  631. [chunk for chunk, _ in batch_data],
  632. [embedding.tolist() for _, embedding in batch_data]
  633. ]
  634. # 插入数据并获取插入操作的结果
  635. insert_result = collection.insert(entities)
  636. # 获取插入的 ID
  637. inserted_ids = insert_result.primary_keys
  638. return inserted_ids
  639. logger.info(f"成功插入{len(batch_data)}个文本块到Milvus")
  640. except Exception as e:
  641. logger.error(f"批量插入数据时发生错误: {str(e)}")
  642. @staticmethod
  643. def _create_index_and_load(collection):
  644. index_params = {
  645. "index_type": "IVF_FLAT",
  646. "metric_type": "L2",
  647. "params": {"nlist": 768}
  648. }
  649. collection.create_index("embedding", index_params)
  650. collection.load()
  651. @staticmethod
  652. def async_analysis(document_id, start_page, end_page, max_tokens):
  653. """
  654. 异步分析文档内容。
  655. :param document_id: 文档的唯一标识符。
  656. :param start_page: 分析开始的页面号。
  657. :param end_page: 分析结束的页面号。
  658. :param max_tokens: 分析的最大令牌数。
  659. 1 已解析、 2 未知: 、 3: 成功、 4: 失败、5: 待 处理
  660. """
  661. start_time = time.time()
  662. excel_status = 1
  663. result = []
  664. logger.info(f"开始处理文档 {document_id}")
  665. try:
  666. DocumentKbm.objects.filter(id=document_id).update(run=1) # 更新文档状态为处理中
  667. document = get_object_or_404(DocumentKbm, id=document_id)
  668. object_name = document.location
  669. file_extension = object_name.split('.')[-1].lower()
  670. knowledgebase = get_object_or_404(Knowledgebase, id=document.kb_id)
  671. bucket_name = knowledgebase.location
  672. logger.info(f"文档信息: object_name:{object_name}, file_extension:{file_extension}, bucket_name:{bucket_name}")
  673. KbmService.clearPreviousData(document_id, bucket_name)
  674. response = minio_client.get_object(bucket_name, object_name)
  675. if not response:
  676. raise ValueError(f"MinIO响应为空。行号:{sys._getframe().f_lineno}")
  677. response_content = response.read()
  678. if not response_content:
  679. raise ValueError(f"从MinIO读取的文件内容为空。行号:{sys._getframe().f_lineno}")
  680. file_content = BytesIO(response_content)
  681. if file_content.getbuffer().nbytes == 0:
  682. raise ValueError(f"文件内容为空。行号:{sys._getframe().f_lineno}")
  683. if file_extension in ['xls', 'xlsx']:
  684. logger.info(f"开始处理Excel文件: {document_id}")
  685. result, excel_status = KbmService.process_excel(file_content, document_id, max_tokens, bucket_name)
  686. elif file_extension == 'pdf':
  687. logger.info(f"开始处理PDF文件: {document_id}")
  688. result = KbmService.process_pdf(file_content, document_id, max_tokens, bucket_name)
  689. elif file_extension == 'md':
  690. logger.info(f"开始处理MD文件: {document_id}")
  691. result = KbmService.process_markdown(file_content, document_id, max_tokens, bucket_name)
  692. elif file_extension in ['doc', 'docx']:
  693. logger.info(f"开始将DOC/DOCX文件转换为PDF: {document_id}")
  694. pdf_content = KbmService.convert_doc_to_pdf(file_content)
  695. result = KbmService.process_pdf(pdf_content, document_id, max_tokens, bucket_name)
  696. if not result:
  697. raise ValueError("处理结果为空。")
  698. KbmService.saveTask(document_id, len(result))
  699. end_time = time.time()
  700. execution_time = round(end_time - start_time, 2)
  701. KbmService.updateDocument(max_tokens, len(result), document_id, execution_time)
  702. if excel_status == 6:
  703. DocumentKbm.objects.filter(id=document_id).update(run=6) # Excel 特殊情况
  704. logger.info(f"文档 {document_id} 更新完成,状态设置为6(Excel特殊情况)")
  705. else:
  706. DocumentKbm.objects.filter(id=document_id).update(run=3) # 假设3表示成功状态
  707. logger.info(f"文档 {document_id} 更新完成,状态设置为3(成功)")
  708. 861
  709. logger.info(f"文档 {document_id} 处理完成")
  710. except Exception as e:
  711. # logger.error(f"处理文档 {document_id} 时发生错误: {str(e)}")
  712. logger.error(f"异常堆栈跟踪:{str(e)}")
  713. traceback.print_exc()
  714. DocumentKbm.objects.filter(id=document_id).update(run=4) # 更新文档状态为失败
  715. @staticmethod
  716. def convert_doc_to_pdf(file_content):
  717. try:
  718. # 创建临时文件
  719. with tempfile.NamedTemporaryFile(delete=False, suffix='.docx') as temp_input:
  720. temp_input.write(file_content.getvalue())
  721. temp_input_path = temp_input.name
  722. temp_output_dir = tempfile.mkdtemp()
  723. # 查找 LibreOffice 路径
  724. libreoffice_path = KbmService.get_libreoffice_path()
  725. if not libreoffice_path:
  726. raise FileNotFoundError("找不到 LibreOffice 可执行文件")
  727. # 转换为 PDF
  728. pdf_path = KbmService.run_libreoffice_conversion(libreoffice_path, temp_input_path, temp_output_dir)
  729. # 读取 PDF 内容
  730. with open(pdf_path, 'rb') as pdf_file:
  731. pdf_content = pdf_file.read()
  732. # 读取并返回 PDF 内容
  733. with open(pdf_path, 'rb') as pdf_file:
  734. return BytesIO(pdf_file.read())
  735. except Exception as e:
  736. logger.error(f"将文档转换为 PDF 时出错: {str(e)}", exc_info=True)
  737. return BytesIO()
  738. finally:
  739. KbmService.cleanup_temp_files(temp_input_path, temp_output_dir)
  740. @staticmethod
  741. def get_libreoffice_path():
  742. system = platform.system()
  743. if system == "Windows":
  744. libreoffice_paths = [r"E:\tool\LibreOffice\program\soffice.exe"]
  745. return next((path for path in libreoffice_paths if os.path.exists(path)), None)
  746. else: # Linux 或 macOS
  747. for path in ['/usr/bin/libreoffice', '/usr/bin/soffice', '/opt/libreoffice/program/soffice']:
  748. if os.path.exists(path):
  749. return path
  750. return shutil.which('libreoffice') or shutil.which('soffice')
  751. @staticmethod
  752. def run_libreoffice_conversion(libreoffice_path, input_path, output_dir):
  753. cmd = [
  754. libreoffice_path,
  755. '--headless',
  756. '--convert-to', 'pdf:writer_pdf_Export:{"PageSize":{"Width":21000,"Height":29700}}',
  757. '--outdir', output_dir,
  758. input_path
  759. ]
  760. try:
  761. env = os.environ.copy()
  762. env['HOME'] = '/mnt/ql_api/tmp' # 设置一个临时的 HOME 目录
  763. env['LC_ALL'] = 'C' # 设置一个标准的语言环境
  764. sleep(0.3)
  765. result = subprocess.run(cmd, check=True, capture_output=True, text=True, timeout=60, env=env)
  766. logger.info(f"LibreOffice 转换输出: {result.stdout}")
  767. pdf_filename = os.path.splitext(os.path.basename(input_path))[0] + '.pdf'
  768. pdf_path = os.path.join(output_dir, pdf_filename)
  769. if not os.path.exists(pdf_path):
  770. raise FileNotFoundError(f"PDF 文件未生成。输出目录内容: {os.listdir(output_dir)}")
  771. # 使用 PyPDF2 检查页数
  772. with open(pdf_path, 'rb') as pdf_file:
  773. pdf_reader = PyPDF2.PdfReader(pdf_file)
  774. page_count = len(pdf_reader.pages)
  775. logger.info(f"生成的 PDF 文件页数: {page_count}")
  776. pdf_size = os.path.getsize(pdf_path)
  777. if pdf_size < 1000:
  778. logger.warning(f"生成的 PDF 文件大小异常小: {pdf_size} bytes")
  779. return pdf_path
  780. except subprocess.TimeoutExpired:
  781. raise TimeoutError("LibreOffice 转换超时")
  782. except subprocess.CalledProcessError as e:
  783. raise RuntimeError(f"LibreOffice 转换失败: {e.output}")
  784. @staticmethod
  785. def cleanup_temp_files(temp_input_path, temp_output_dir):
  786. if os.path.exists(temp_input_path):
  787. os.remove(temp_input_path)
  788. if os.path.exists(temp_output_dir):
  789. shutil.rmtree(temp_output_dir)
  790. @staticmethod
  791. def process_excel(file_content, document_id, max_tokens, bucket_name):
  792. result = []
  793. collection_name = f"{bucket_name}"
  794. status=1
  795. # object1 object1为后续可能添加的字段 因为无法直接修改名称 备用
  796. source = "知识库"
  797. object1 = "some_object1"
  798. object2 = "some_object2"
  799. def warning_catcher(message, category, filename, lineno, file=None, line=None):
  800. nonlocal status
  801. if category == UserWarning:
  802. if "File contains an invalid specification for 0" in str(message) or \
  803. "Defined names for sheet index 0 cannot be located" in str(message):
  804. status = 6
  805. logger.error(f"Warning: {message}")
  806. warnings.showwarning = warning_catcher
  807. try:
  808. excel_file = pd.ExcelFile(file_content)
  809. except Exception as e:
  810. logger.error(f"Error reading Excel file: {str(e)}")
  811. return result
  812. try:
  813. connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT,user=MILVUS_USER,password=MILVUS_PASSWORD)
  814. collection = KbmService._get_or_create_collection(collection_name)
  815. for sheet_name in excel_file.sheet_names:
  816. df = pd.read_excel(excel_file, sheet_name=sheet_name)
  817. if df.empty:
  818. logger.warning(f"Sheet '{sheet_name}' is empty. Skipping.")
  819. continue
  820. logger.info(f"Processing sheet '{sheet_name}' with shape {df.shape}")
  821. markdown_content = KbmService._excel_to_markdown(df, sheet_name)
  822. chunks = KbmService._split_markdown(markdown_content, max_tokens)
  823. for chunk_number, chunk_content in enumerate(chunks, start=1):
  824. try:
  825. if not chunk_content.strip():
  826. logger.warning(f"Empty chunk {chunk_number} in sheet '{sheet_name}'. Skipping.")
  827. continue
  828. embedding = KbmService.get_embedding_excel(chunk_content, target_dim=VECTOR_DIMENSION)
  829. if isinstance(embedding, (list, np.ndarray)) and len(embedding) == VECTOR_DIMENSION:
  830. sleep(1)
  831. logger.info("减少milvus压力睡眠1秒")
  832. milvus_id = KbmService._insert_data(collection, chunk_content, embedding,source,object1,object2)
  833. else:
  834. logger.error(f"Invalid embedding format for chunk {chunk_number} of sheet {sheet_name}.")
  835. continue
  836. KbmService.saveTaskSublist(
  837. document_id=document_id,
  838. name=f"sheet_{sheet_name}",
  839. page_number=1,
  840. chunk_number=chunk_number,
  841. content=chunk_content,
  842. milvus_id=milvus_id
  843. )
  844. result.append({
  845. 'page_number': 1,
  846. 'chunk_number': chunk_number,
  847. })
  848. except Exception as e:
  849. logger.error(f"Error processing chunk {chunk_number} of sheet {sheet_name}: {str(e)}")
  850. logger.error(f"Chunk content: {chunk_content}")
  851. logger.exception("Detailed error information:")
  852. except Exception as e:
  853. logger.error(f"Error processing Excel file: {str(e)}")
  854. logger.error("Detailed error information:")
  855. logger.error(traceback.format_exc())
  856. status = 6 # 设置状态为6,表示处理出错
  857. raise
  858. finally:
  859. connections.disconnect("default")
  860. return result, status
  861. #excel转markdown
  862. @staticmethod
  863. def _excel_to_markdown(df, sheet_name):
  864. if df.empty:
  865. return f"# {sheet_name}\n\n表格为空"
  866. headers = df.columns.tolist()
  867. data = df.values.tolist()
  868. # 将所有数据转换为字符串
  869. data = [[str(cell) for cell in row] for row in data]
  870. markdown = f"# {sheet_name}\n\n"
  871. markdown += tabulate(data, headers=headers, tablefmt="pipe", showindex=False)
  872. return markdown
  873. #excel分割策略
  874. @staticmethod
  875. def _split_json(json_str, max_tokens):
  876. # 简单的分割策略,可以根据需要优化
  877. data = json.loads(json_str)
  878. chunks = []
  879. current_chunk = []
  880. current_size = 0
  881. for item in data:
  882. item_str = json.dumps(item)
  883. item_size = len(item_str)
  884. if current_size + item_size > max_tokens and current_chunk:
  885. chunks.append(json.dumps(current_chunk))
  886. current_chunk = []
  887. current_size = 0
  888. current_chunk.append(item)
  889. current_size += item_size
  890. if current_chunk:
  891. chunks.append(json.dumps(current_chunk))
  892. return chunks
  893. @staticmethod
  894. def _split_markdown(markdown_content, max_tokens):
  895. chunks = []
  896. current_chunk = ""
  897. lines = markdown_content.split('\n')
  898. for line in lines:
  899. if len(current_chunk) + len(line) + 1 > max_tokens:
  900. if current_chunk:
  901. chunks.append(current_chunk.strip())
  902. current_chunk = line
  903. else:
  904. current_chunk += '\n' + line if current_chunk else line
  905. if current_chunk:
  906. chunks.append(current_chunk.strip())
  907. return chunks
  908. #milvus excel插入格式
  909. @staticmethod
  910. def _insert_data(collection, content, embedding,source,object1,object2):
  911. try:
  912. data = [
  913. [source], # content field
  914. [object1], # content field
  915. [object2], # content field
  916. [content], # content field
  917. [embedding] # embedding field
  918. ]
  919. insert_result = collection.insert(data)
  920. logger.info(f"Inserted 1 record into Milvus")
  921. return insert_result.primary_keys[0] # 返回插入的 ID
  922. except Exception as e:
  923. logger.error(f"Error inserting data into Milvus: {str(e)}")
  924. raise
  925. #创建milvus索引
  926. @staticmethod
  927. def _create_index_if_not_exists(collection):
  928. if not collection.has_index():
  929. index_params = {
  930. "index_type": "IVF_FLAT",
  931. "metric_type": "L2",
  932. "params": {"nlist": 768}
  933. }
  934. collection.create_index("embedding", index_params)
  935. @staticmethod
  936. def process_pdf(file_content, document_id, max_tokens, bucket_name):
  937. """
  938. 解析PDF文件并处理每一页的内容。
  939. Args:
  940. file_content (BytesIO): PDF文件的内容。
  941. document_id (str): 文档的ID。
  942. max_tokens (int): 最大令牌数。
  943. bucket_name (str): 存储桶的名称。
  944. Returns:
  945. list: 处理后的页面列表。
  946. """
  947. logger.info(f"开始处理 PDF,document_id: {document_id}, max_tokens: {max_tokens}")
  948. text_chunks = []
  949. try:
  950. sentence_chunks, page_images = pdf_utils.PDFProcessor(2).get_results(file_content, document_id, max_tokens, bucket_name, KbmService)
  951. # 一张页面 , 一张页面的完这
  952. for sentence_chunk in sentence_chunks:
  953. text_chunks.append( KbmService.split_text_by_semantic(sentence_chunk, max_tokens, bucket_name))
  954. sleep(.5)
  955. logger.info("分割完成")
  956. result = []
  957. inder = 1
  958. for i, chunks in enumerate(text_chunks,1):
  959. image_name = page_images[i-1][0]
  960. for j, chunk in enumerate(chunks,1):
  961. KbmService.saveTaskSublist(
  962. document_id=document_id,
  963. name=image_name,
  964. page_number=i,
  965. chunk_number=inder,
  966. content=chunk['content'],
  967. milvus_id=chunk['milvus_id']
  968. )
  969. inder +=1
  970. result.append({
  971. 'page_number': i,
  972. 'chunk_number': len(chunks),
  973. })
  974. logger.info("解析结束")
  975. return result
  976. except Exception as e:
  977. logger.error(f'错误 {str(e)}')
  978. # raise ValueError(f'{str(e)} 行号:{sys._getframe().f_lineno}')
  979. return []
  980. #解析markdown
  981. @staticmethod
  982. def process_markdown(file_content, document_id, max_tokens, bucket_name):
  983. logger.info(f"开始解析 Markdown,document_id: {document_id}, max_tokens: {max_tokens}")
  984. try:
  985. # 检测文件编码
  986. raw_content = file_content.read()
  987. detected = chardet.detect(raw_content)
  988. encoding = detected['encoding']
  989. logger.info(f"检测到的文件编码: {encoding}")
  990. # 解码文件内容
  991. text = raw_content.decode(encoding)
  992. logger.info(f"Markdown 文件总字符数: {len(text)}")
  993. logger.debug(f"Markdown 文件前100个字符: {text[:100]}")
  994. # 将 Markdown 转换为 HTML
  995. html = markdown.markdown(text)
  996. # 使用 BeautifulSoup 提取纯文本
  997. soup = BeautifulSoup(html, 'html.parser')
  998. plain_text = soup.get_text()
  999. logger.info(f"提取的纯文本总字符数: {len(plain_text)}")
  1000. logger.debug(f"提取的纯文本前100个字符: {plain_text[:100]}")
  1001. if not plain_text.strip():
  1002. logger.warning("Markdown 文件内容为空")
  1003. return []
  1004. sentence_chunks = pdf_utils.PDFProcessor(1).split_sentences(plain_text)
  1005. # 使用 split_text_by_semantic 方法分割文本
  1006. text_chunks = KbmService.split_text_by_semantic(sentence_chunks, max_tokens, bucket_name)
  1007. logger.info(f"分割后的文本块数: {len(text_chunks)}")
  1008. result = []
  1009. for i, chunk in enumerate(text_chunks, 1):
  1010. KbmService.saveTaskSublist(
  1011. document_id=document_id,
  1012. name="markdown_content",
  1013. page_number=1,
  1014. chunk_number=i,
  1015. content=chunk['content'],
  1016. milvus_id=chunk['milvus_id']
  1017. )
  1018. result.append({
  1019. 'page_number': 1,
  1020. 'chunk_number': i,
  1021. })
  1022. logger.info(f"Markdown 处理完成,总共生成 {len(result)} 个文本块")
  1023. return result
  1024. except Exception as e:
  1025. logger.error(f"处理 Markdown 时发生错误: {str(e)}")
  1026. logger.exception("详细错误信息:")
  1027. return []
  1028. @staticmethod
  1029. def split_text(text, max_tokens):
  1030. words = text.split()
  1031. chunks = []
  1032. current_chunk = []
  1033. current_token_count = 0
  1034. for word in words:
  1035. word_tokens = KbmService.estimate_tokens(word)
  1036. if current_token_count + word_tokens > max_tokens and current_chunk:
  1037. chunks.append(' '.join(current_chunk))
  1038. current_chunk = []
  1039. current_token_count = 0
  1040. current_chunk.append(word)
  1041. current_token_count += word_tokens
  1042. if current_chunk:
  1043. chunks.append(' '.join(current_chunk))
  1044. return chunks
  1045. @staticmethod
  1046. def estimate_tokens(text):
  1047. return len(re.findall(r'\w+', text)) * 1.3
  1048. @staticmethod
  1049. def extract_text_from_image(image_data):
  1050. try:
  1051. if isinstance(image_data, BytesIO):
  1052. image_data = image_data.getvalue()
  1053. # 读取图像
  1054. nparr = np.frombuffer(image_data, np.uint8)
  1055. image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
  1056. # 图像预处理
  1057. preprocessed = KbmService.preprocess_image(image)
  1058. # 使用进行识别
  1059. # result = Paddle_OCR.ocr(preprocessed, cls=True)
  1060. result = ocr_utils.OCRProcesso().ocr(preprocessed)
  1061. # 后处理
  1062. text = KbmService.post_process_text(result)
  1063. if text and len(text) > 2: # 假设有意义的文本至少有3个字符
  1064. logger.info(f"提取的文本长度: {len(text)} {text}")
  1065. return text
  1066. else:
  1067. logger.info("提取的内容似乎是图像,而不是文本")
  1068. return "图片"
  1069. except Exception as e:
  1070. logger.error(f"从图像提取文本时出错: {str(e)} {sys._getframe().f_lineno}")
  1071. return "图片"
  1072. @staticmethod
  1073. def preprocess_image(image):
  1074. # 转换为灰度图像
  1075. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  1076. # 自适应阈值处理
  1077. binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
  1078. # 对二值图像进行膨胀操作,使文字更粗
  1079. kernel = np.ones((2, 2), np.uint8)
  1080. dilated = cv2.dilate(binary, kernel, iterations=1)
  1081. return dilated
  1082. def post_process_text(text):
  1083. if text:
  1084. # 将连续的冒号或点替换为空格
  1085. text = re.sub(r'[:.]+', ' ', text)
  1086. # 保留中文字符、英文字母、数字、常用标点
  1087. text = re.sub(r'[^\u4e00-\u9fff\u3000-\u303fa-zA-Z0-9.,!?;:()"\'\s]', '', text)
  1088. # 删除连续的数字(3个或更多)
  1089. text = re.sub(r'\d{3,}', '', text)
  1090. # 处理多余的空白字符
  1091. text = re.sub(r'\s+', ' ', text).strip()
  1092. # 删除单独的数字,但保留章节编号和有意义的数字
  1093. text = re.sub(r'\b(?<![第章])\d+(?!\d)\b', '', text)
  1094. # 清理多余的空格
  1095. text = re.sub(r'\s+', ' ', text).strip()
  1096. return text
  1097. @staticmethod
  1098. def render_page_to_image(page, scale=2):
  1099. # logger.info(page)
  1100. try:
  1101. pix = page.get_pixmap(matrix=fitz.Matrix(scale, scale))
  1102. img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
  1103. buffered = BytesIO()
  1104. img.save(buffered, format="PNG")
  1105. # image_name = rf"E:/page_image_{uuid.uuid4()}.png"
  1106. # img.save(image_name)
  1107. buffered.seek(0)
  1108. return buffered
  1109. except Exception as e:
  1110. logger.error('文件数据不正确')
  1111. raise ValueError("文件数据不正确")
  1112. @staticmethod
  1113. def save_image_to_minio(image_data, bucket_name):
  1114. image_name = f"page_image_{uuid.uuid4()}.png"
  1115. minio_client.put_object(bucket_name, image_name, image_data, length=image_data.getbuffer().nbytes)
  1116. return image_name
  1117. @staticmethod
  1118. @transaction.atomic
  1119. def saveTask(document_id, total_chunks):
  1120. Task.objects.update_or_create(
  1121. doc_id=document_id, # 查找条件
  1122. defaults={'to_page': total_chunks} # 要更新或创建的字段
  1123. )
  1124. @staticmethod
  1125. @transaction.atomic
  1126. def saveTaskSublist(document_id, name, page_number=None, chunk_number=None, content=None,milvus_id=None):
  1127. try:
  1128. # 确保 content 是 Unicode 字符串
  1129. if content is not None:
  1130. if isinstance(content, bytes):
  1131. content = content.decode('utf-8')
  1132. else:
  1133. content = str(content)
  1134. TaskSublist.objects.create(
  1135. doc_id=document_id,
  1136. name=name,
  1137. page_number=str(page_number) if page_number is not None else '0',
  1138. chunk_number=str(chunk_number) if chunk_number is not None else '0',
  1139. content=content,
  1140. milvus_id=milvus_id
  1141. )
  1142. logger.info(f"Successfully saved TaskSublist for document {document_id}, chunk {chunk_number}")
  1143. except Exception as e:
  1144. logger.error(f"Error saving TaskSublist: {str(e)}")
  1145. # 可以选择在这里重新抛出异常,或者进行其他错误处理
  1146. @staticmethod
  1147. @transaction.atomic
  1148. def clearPreviousData(document_id,bucket_name):
  1149. try:
  1150. # 获取与文档相关的所有 TaskSublist,milvus 记录
  1151. task_sublists = TaskSublist.objects.filter(doc_id=document_id)
  1152. # 获取集合对象
  1153. connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT,user=MILVUS_USER,password=MILVUS_PASSWORD)
  1154. milvus_collection_exists = utility.has_collection(bucket_name)
  1155. if milvus_collection_exists:
  1156. collection = Collection(bucket_name)
  1157. # 连接到 Milvus
  1158. # 从 MinIO 中删除相关的图片
  1159. for task in task_sublists:
  1160. try:
  1161. # 执行删除minio
  1162. minio_client.remove_object(bucket_name, task.name)
  1163. # 执行删除milvus
  1164. if milvus_collection_exists and task.milvus_id:
  1165. expr = f'id in [{task.milvus_id}]'
  1166. collection.delete(expr)
  1167. except Exception as e:
  1168. logger.error(f"Error deleting object from MinIO: {e}")
  1169. # 从数据库中删除 TaskSublist 记录
  1170. task_sublists.delete()
  1171. except Exception as e:
  1172. logger.error(f"Error deleting object from MinIO: {str(e)} {sys._getframe().f_lineno}")
  1173. finally:
  1174. connections.disconnect("default")
  1175. @staticmethod
  1176. @transaction.atomic
  1177. def updateDocument(max_tokens, total_chunks, document_id,execution_time):
  1178. try:
  1179. # 检查是否存在相关的 TaskSublist
  1180. count = TaskSublist.objects.filter(doc_id=document_id).count()
  1181. # 根据 TaskSublist 的存在与否设置运行状态
  1182. progress_status = 1 if count > 0 else -1
  1183. # 更新 DocumentKbm 对象
  1184. updated = DocumentKbm.objects.filter(id=document_id).update(
  1185. token_num=max_tokens,
  1186. chunk_num=total_chunks,
  1187. progress=progress_status,
  1188. process_begin_at=timezone.now(),
  1189. process_duation= execution_time
  1190. )
  1191. if updated:
  1192. return True, f"Document {document_id} updated successfully."
  1193. else:
  1194. return False, f"Document {document_id} not found."
  1195. except Exception as e:
  1196. # 如果发生任何错误,事务会自动回滚
  1197. return False, f"Error updating document: {str(e)}"
  1198. #异步调用
  1199. @staticmethod
  1200. def searchTaskInfo(request):
  1201. document_id = request.POST.get("document_id")
  1202. page = request.POST.get('page', 1)
  1203. page_size = request.POST.get('page_size', 10) # 每页显示的项目数,默认为10
  1204. taskSublists = TaskSublist.objects.filter(doc_id=document_id).order_by('id')
  1205. document = get_object_or_404(DocumentKbm, id=document_id)
  1206. location = document.location
  1207. knowledgebase = get_object_or_404(Knowledgebase, id=document.kb_id)
  1208. bucket_name = knowledgebase.location
  1209. documentUrl = minio_client.presigned_get_object(
  1210. bucket_name=bucket_name,
  1211. object_name=location,
  1212. expires=timedelta(days=1) # URL有效期为1天
  1213. )
  1214. # 创建分页器
  1215. paginator = Paginator(taskSublists, page_size)
  1216. try:
  1217. tasks_page = paginator.page(page)
  1218. except PageNotAnInteger:
  1219. # 如果页码不是整数,返回第一页
  1220. tasks_page = paginator.page(1)
  1221. except EmptyPage:
  1222. # 如果页码超出范围,返回最后一页
  1223. tasks_page = paginator.page(paginator.num_pages)
  1224. task_results = []
  1225. for task in tasks_page:
  1226. try:
  1227. # 生成MinIO对象的预签名URL
  1228. url = minio_client.presigned_get_object(
  1229. bucket_name=bucket_name,
  1230. object_name=task.name,
  1231. expires=timedelta(days=1) # URL有效期为1天
  1232. )
  1233. task_results.append({
  1234. 'id': task.id,
  1235. 'doc_id': task.doc_id,
  1236. 'name': task.name,
  1237. 'page_number': task.page_number,
  1238. 'chunk_number': task.chunk_number,
  1239. 'content': task.content,
  1240. 'url': url
  1241. })
  1242. except Exception as e:
  1243. logger.error(f"Error generating URL for object {task.name}: {str(e)}")
  1244. # 如果生成URL失败,我们仍然添加其他信息,但URL为None
  1245. task_results.append({
  1246. 'id': task.id,
  1247. 'doc_id': task.doc_id,
  1248. 'name': task.name,
  1249. 'page_number': task.page_number,
  1250. 'chunk_number': task.chunk_number,
  1251. 'content': task.content,
  1252. 'url': None
  1253. })
  1254. # 创建包含 documentUrl 和分页信息的最终结果
  1255. result = {
  1256. 'documentUrl': documentUrl,
  1257. 'tasks': task_results,
  1258. 'pagination': {
  1259. 'current_page': tasks_page.number,
  1260. 'num_pages': paginator.num_pages,
  1261. 'per_page': page_size,
  1262. 'total_count': paginator.count,
  1263. 'has_next': tasks_page.has_next(),
  1264. 'has_previous': tasks_page.has_previous(),
  1265. }
  1266. }
  1267. return success(result)
  1268. @staticmethod
  1269. @transaction.atomic
  1270. def deleteBucket(request):
  1271. bucket_id = request.POST.get("bucket_id")
  1272. if not bucket_id:
  1273. return fail("Bucket ID 为空")
  1274. try:
  1275. # 检查是否存在未删除的文档
  1276. active_docs_count = DocumentKbm.objects.filter(kb_id=bucket_id).exclude(status=4).count()
  1277. if active_docs_count > 0:
  1278. return fail(f"无法删除知识库,还有 {active_docs_count} 个未删除的文档")
  1279. # 如果没有未删除的文档,则更新知识库状态
  1280. updated_count = Knowledgebase.objects.filter(id=bucket_id).update(status=4,name=bucket_id, location=bucket_id)
  1281. if updated_count == 0:
  1282. return fail("指定的知识库不存在")
  1283. return success("知识库已成功删除")
  1284. except Exception as e:
  1285. return fail(f"删除知识库时发生错误: {str(e)}")
  1286. @staticmethod
  1287. def getRunStatus(request):
  1288. document_id = request.POST.get("document_id")
  1289. run = DocumentKbm.objects.filter(id = document_id).values("run").first()
  1290. return success(run)
  1291. @staticmethod
  1292. def batchAnalysis(request):
  1293. ids_str = request.POST.get("ids")
  1294. start_page = int(request.POST.get('start_page', 1))
  1295. end_page = int(request.POST.get('end_page', -1))
  1296. max_tokens = int(request.POST.get('max_tokens', 2048))
  1297. try:
  1298. # 尝试将字符串解析为 JSON 列表
  1299. ids = json.loads(ids_str)
  1300. if not isinstance(ids, list):
  1301. return fail("无效输入:'ids'应该是一个列表")
  1302. results = []
  1303. for document_id in ids:
  1304. sleep(0.1)
  1305. logger.info("缓解压力沉睡0.1秒")
  1306. # 为每个 document_id 创建一个新的请求对象
  1307. analysis_request = type('AnalysisRequest', (), {})()
  1308. analysis_request.POST = {
  1309. 'document_id': document_id,
  1310. 'start_page': start_page,
  1311. 'end_page': end_page,
  1312. 'max_tokens': max_tokens
  1313. }
  1314. # 调用 analysis 方法
  1315. response = KbmService.analysis(analysis_request)
  1316. message = response.get('message')
  1317. return success(message,"已添加到队列")
  1318. except Exception as e:
  1319. return fail(f"An error occurred: {str(e)}")
  1320. # 假设这是您支持的文件后缀名列表
  1321. SUPPORTED_SUFFIXES = [
  1322. 'txt', 'pdf', 'doc', 'docx', 'xls', 'xlsx', 'md'
  1323. ]
  1324. @staticmethod
  1325. def getSuffixName(request):
  1326. try:
  1327. # 获取数据库中的所有不重复的 type 值
  1328. db_types = DocumentKbm.objects.values_list('type', flat=True).distinct()
  1329. # 将数据库中的类型转换为集合
  1330. db_types_set = set(db_types)
  1331. # 将 SUPPORTED_SUFFIXES 转换为集合
  1332. supported_set = set(KbmService.SUPPORTED_SUFFIXES)
  1333. # 合并两个集合,自动去除重复项
  1334. combined_set = supported_set.union(db_types_set)
  1335. # 将结果转换回列表
  1336. combined_suffixes = list(combined_set)
  1337. # 对结果进行排序(可选)
  1338. combined_suffixes.sort()
  1339. return success(combined_suffixes)
  1340. except Exception as e:
  1341. return fail(f"获取文件后缀名时发生错误: {str(e)}")
  1342. @staticmethod
  1343. @transaction.atomic
  1344. def batchMove(request):
  1345. ids = json.loads(request.POST.get("ids"))
  1346. doc_type_id = request.POST.get("doc_type_id")
  1347. if not doc_type_id:
  1348. return fail("分类id为空")
  1349. if not ids:
  1350. return fail("未传出文件id")
  1351. type = KbmDocumentType.objects.filter(id=doc_type_id).exclude(status=4).first()
  1352. if not type:
  1353. return fail("当前分类不存在")
  1354. DocumentKbm.objects.filter(id__in=ids).update(doc_type_id=doc_type_id)
  1355. return success("批量移动成功")
  1356. @staticmethod
  1357. @transaction.atomic
  1358. def moveDocument(doc_id, doc_type_id):
  1359. # 这个方法可以保持不变,因为它已经是单次更新操作
  1360. return DocumentKbm.objects.filter(id=doc_id).update(doc_type_id=doc_type_id)
  1361. @staticmethod
  1362. @transaction.atomic
  1363. def updateKbm(request):
  1364. """
  1365. 更新知识库信息
  1366. Args:
  1367. request (HttpRequest): 请求对象
  1368. Returns:
  1369. dict: 成功或失败的响应字典
  1370. """
  1371. id = request.POST.get("id")
  1372. if not id:
  1373. return fail("id为空")
  1374. kmb = Knowledgebase.objects.filter(id=id).first()
  1375. name = request.POST.get("name")
  1376. if not name:
  1377. return fail("名称不能为空")
  1378. if kmb:
  1379. kmb.name = name
  1380. kmb.description = request.POST.get("description","")
  1381. kmb.save()
  1382. return success("修改成功")
  1383. else:
  1384. return fail("修改失败")