chroma新增、删除、知识库应用
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

112 lines
4.6 KiB

  1. # coding:utf8
  2. import os, sys
  3. import io
  4. from jsonpath_ng import jsonpath, parse
  5. import uuid
  6. sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf8')
  7. cur_dir = os.path.dirname(os.path.abspath(__file__)) or os.getcwd()
  8. par_dir = os.path.abspath(os.path.join(cur_dir, os.path.pardir))
  9. sys.path.append(cur_dir)
  10. sys.path.append(par_dir)
  11. import json
  12. from django.http import HttpResponse
  13. from text_analysis.tools import to_kafka
  14. from django.views.decorators.csrf import csrf_exempt
  15. from log_util.set_logger import set_logger
  16. logging = set_logger('logs/results.log')
  17. import traceback
  18. import queue
  19. import requests
  20. from text_analysis.tools.tool import parse_data
  21. from text_analysis.chroma1 import LangChainChroma
  22. import time
  23. from datetime import datetime
  24. import os
  25. # 任务队列
  26. global task_queue
  27. task_queue = queue.Queue()
  28. # LC = LangChainChroma()
  29. @csrf_exempt
  30. def createChroma(request):
  31. if request.method == 'POST':
  32. try:
  33. raw_data = json.loads(request.body)
  34. task_queue.put(raw_data)
  35. return HttpResponse(json.dumps({"code": 1, "msg": "请求正常!"}, ensure_ascii=False))
  36. except:
  37. logging.error(traceback.format_exc())
  38. return HttpResponse(json.dumps({"code": 0, "msg": "请求json格式不正确!"}, ensure_ascii=False))
  39. else:
  40. return HttpResponse(json.dumps({"code": 0, "msg": "请求方式错误,改为post请求"}, ensure_ascii=False))
  41. def upload():
  42. while True:
  43. try:
  44. if task_queue.qsize() > 0:
  45. logging.info("取任务队列长度{}".format(task_queue.qsize()))
  46. raw_data = task_queue.get()
  47. output=raw_data["output"]
  48. res_tmp={key: "" for key in output}
  49. if "id" in res_tmp.keys():
  50. res_tmp["id"]=str(uuid.uuid4())
  51. res_tmp["isLast"]=1
  52. logging.info("任务数据为:{}".format(raw_data))
  53. # chunkSize=parse_data(raw_data,raw_data["input"]["chunkSize"])
  54. content=parse_data(raw_data,raw_data["input"]["content"])
  55. fieldName=parse_data(raw_data,raw_data["input"]["fieldName"])
  56. dataId=parse_data(raw_data,raw_data["input"]["dataId"])
  57. # dataId=raw_data["dataId"]
  58. if content and fieldName and dataId:
  59. vector_db=LangChainChroma(fieldName)
  60. docs=vector_db.text_splitter.split_text(content)
  61. res,db_count=vector_db.add_documents(docs,dataId)
  62. vector_db.db_close()
  63. logging.info('当前数据划分{}个块。数据库{}共有{}个块'.format(len(res), fieldName,db_count))
  64. # res=LC.addChroma(content,fieldName,logging,chunkSize)
  65. res_tmp['resultsID']=res
  66. raw_data["result"] = {"successCode": "", "errorLog": "", "results": ""}
  67. if res:
  68. res_tmp["status"]=1
  69. raw_data["result"]["successCode"] = "1"
  70. raw_data["result"]["status"] = 1
  71. raw_data["result"]["message"] = "成功"
  72. else:
  73. res_tmp["status"]=3
  74. raw_data["result"]["successCode"] = "0"
  75. raw_data["result"]["status"] = 2
  76. raw_data["result"]["message"] = "异常"
  77. else:
  78. res_tmp["status"] = 3
  79. raw_data["result"]["successCode"] = "0"
  80. raw_data["result"]["errorLog"] = "请检查content/fieldName/dataId,要求非空"
  81. raw_data["result"]["status"] = 2
  82. raw_data["result"]["message"] = "请检查content/fieldName/dataId,要求非空"
  83. res_tmp_json = json.dumps(res_tmp, ensure_ascii=False)
  84. raw_data["result"]["results"]=res_tmp_json
  85. logging.info("结果数据为:{}".format(raw_data))
  86. to_kafka.send_kafka(raw_data, logging)
  87. else:
  88. # 暂无任务,进入休眠
  89. time.sleep(10)
  90. except:
  91. raw_data["result"]={}
  92. raw_data["result"]["successCode"] = "0"
  93. raw_data["result"]["errorLog"] = traceback.format_exc()
  94. res_tmp["status"] = 3
  95. raw_data["result"]["status"] = 2
  96. raw_data["result"]["message"] = "异常"
  97. raw_data["result"]["results"] = json.dumps(res_tmp, ensure_ascii=False)
  98. logging.error(traceback.format_exc())
  99. to_kafka.send_kafka(raw_data, logging)