KbmService.py 65 KB

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