rabbitmq_Process.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104
  1. # 新rabbitmq队列
  2. # 导入 backend/__init__.py 中的所有内容
  3. from .... import *
  4. import pika
  5. from DCbackend import settings
  6. from DCbackend.utils.common import success, fail
  7. from backend.models import Knowledgebase, DocumentKbm, File2document, File, Task, TaskSublist, KbmDocumentType
  8. def send_to_rabbitmq(queue_name, message):
  9. """
  10. 将消息发送到指定的RabbitMQ队列
  11. """
  12. try:
  13. connection = pika.BlockingConnection(pika.ConnectionParameters(
  14. host=settings.RABBITMQ_HOST,
  15. port=settings.RABBITMQ_PORT,
  16. credentials=pika.PlainCredentials(
  17. settings.RABBITMQ_USER,
  18. settings.RABBITMQ_PASSWORD
  19. )
  20. ))
  21. channel = connection.channel()
  22. channel.queue_declare(queue=queue_name, durable=True)
  23. channel.basic_publish(
  24. exchange='',
  25. routing_key=queue_name,
  26. body=json.dumps(message),
  27. properties=pika.BasicProperties(
  28. delivery_mode=2, # 使消息持久化
  29. )
  30. )
  31. connection.close()
  32. logger.info(f"消息已发送到队列 {queue_name}")
  33. connection = None
  34. channel = None
  35. return True
  36. except Exception as e:
  37. logger.error(f"发送消息到RabbitMQ时出错: {str(e)}")
  38. return False
  39. def analysis(request,KbmService):
  40. """
  41. 分析请求并处理RabbitMQ队列中的消息。
  42. Args:
  43. request (object): 需要分析的请求对象。
  44. KbmService (object): KbmService对象,用于发送消息到RabbitMQ队列。
  45. Returns:
  46. None
  47. """
  48. document_id = request.POST.get("document_id")
  49. start_page = int(request.POST.get('start_page', 1))
  50. end_page = int(request.POST.get('end_page', -1))
  51. max_tokens = int(request.POST.get('max_tokens', 2048))
  52. if max_tokens == 0:
  53. max_tokens = 2048
  54. logger.info(f"开始处理文档 ID: {document_id}")
  55. try:
  56. document = DocumentKbm.objects.get(id=document_id)
  57. if int(document.run) in [1, 5]: # 1: 处理中, 5: 等待处理
  58. logger.info(f"文档 {document_id} 已有队列")
  59. return success("文档正在处理中或已经处理完成")
  60. # 准备消息
  61. message = {
  62. 'document_id': document_id,
  63. 'start_page': start_page,
  64. 'end_page': end_page,
  65. 'max_tokens': max_tokens
  66. }
  67. # 发送消息到队列
  68. if KbmService.send_to_rabbitmq(settings.RABBITMQ_QUEUE_NAME, message):
  69. # 更新文档状态为等待处理
  70. document.run = 5 # 5表示等待处理
  71. document.save()
  72. logger.info(f"文档 {document_id} 状态已更新为等待处理")
  73. return success("文档已添加到处理队列")
  74. else:
  75. document.run = 4
  76. document.save()
  77. return fail("添加文档到处理队列失败")
  78. except DocumentKbm.DoesNotExist:
  79. logger.error(f"文档 {document_id} 不存在")
  80. document.run = 4
  81. document.save()
  82. return fail("文档不存在")
  83. except Exception as e:
  84. logger.error(f"处理文档 {document_id} 时出错: {str(e)}")
  85. document.run = 4
  86. document.save()
  87. return fail("处理文档时出错")