查询知识库应用
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.

93 lines
3.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. from datetime import datetime, timedelta
  17. logging = set_logger('logs/results.log')
  18. import traceback
  19. import queue
  20. import requests
  21. from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
  22. from text_analysis.tools.tool import parse_data,promptPro
  23. from text_analysis.chroma1 import LangChainChroma
  24. import time
  25. from datetime import datetime
  26. import os
  27. from langchain.vectorstores import Chroma
  28. # 任务队列
  29. global task_queue
  30. task_queue = queue.Queue()
  31. # LC = LangChainChroma()
  32. @csrf_exempt
  33. def promptSim(request):
  34. if request.method == 'POST':
  35. try:
  36. raw_data = json.loads(request.body)
  37. task_queue.put(raw_data)
  38. return HttpResponse(json.dumps({"code": 1, "msg": "请求正常!"}, ensure_ascii=False))
  39. except:
  40. logging.error(traceback.format_exc())
  41. return HttpResponse(json.dumps({"code": 0, "msg": "请求json格式不正确!"}, ensure_ascii=False))
  42. else:
  43. return HttpResponse(json.dumps({"code": 0, "msg": "请求方式错误,改为post请求"}, ensure_ascii=False))
  44. def upload():
  45. while True:
  46. try:
  47. if task_queue.qsize() > 0:
  48. logging.info("取任务队列长度{}".format(task_queue.qsize()))
  49. raw_data = task_queue.get()
  50. output=raw_data["output"]
  51. res_tmp={key: "" for key in output}
  52. if "id" in res_tmp.keys():
  53. res_tmp["id"]=str(uuid.uuid4())
  54. logging.info("任务数据为:{}".format(raw_data))
  55. topn=parse_data(raw_data,raw_data["input"]["topn"])
  56. prompt=parse_data(raw_data,raw_data["input"]["prompt"])
  57. fieldName=parse_data(raw_data,raw_data["input"]["fieldName"])
  58. vector_db=LangChainChroma(fieldName)
  59. docs=vector_db.search(prompt,topn)
  60. vector_db.db_close()
  61. logging.info("向量数据库搜索的相似上下文:{}".format(docs))
  62. #组装,最长字符5W
  63. res=promptPro(prompt,docs)
  64. logging.info("生成的上下文:{}".format(res))
  65. res_tmp["promptRes"]=res
  66. res_tmp_json = json.dumps(res_tmp, ensure_ascii=False)
  67. raw_data["result"] = {"successCode": "1", "errorLog": "", "results": res_tmp_json}
  68. raw_data["result"]["status"] = 1
  69. raw_data["result"]["message"] = "成功"
  70. logging.info("结果数据为:{}".format(raw_data))
  71. to_kafka.send_kafka(raw_data, logging)
  72. else:
  73. # 暂无任务,进入休眠
  74. time.sleep(10)
  75. except:
  76. raw_data["result"]={}
  77. raw_data["result"]["successCode"] = "0"
  78. raw_data["result"]["errorLog"] = traceback.format_exc()
  79. raw_data["result"]["status"] = 2
  80. raw_data["result"]["message"] = "异常"
  81. raw_data["result"]["results"] = json.dumps(res_tmp, ensure_ascii=False)
  82. logging.error(traceback.format_exc())
  83. to_kafka.send_kafka(raw_data, logging)