factory/apps/ichat/views2.py

129 lines
5.1 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import requests
import os
from apps.utils.sql import execute_raw_sql
import json
from apps.utils.tools import MyJSONEncoder
from .utils import is_safe_sql
from rest_framework.views import APIView
from drf_yasg.utils import swagger_auto_schema
from rest_framework import serializers
from rest_framework.exceptions import ParseError
from rest_framework.response import Response
from django.conf import settings
from apps.utils.mixins import MyLoggingMixin
from django.core.cache import cache
import uuid
from apps.utils.thread import MyThread
LLM_URL = getattr(settings, "LLM_URL", "")
API_KEY = getattr(settings, "LLM_API_KEY", "")
MODEL = "qwen14b"
HEADERS = {
"Authorization": f"Bearer {API_KEY}",
"Content-Type": "application/json"
}
CUR_DIR = os.path.dirname(os.path.abspath(__file__))
def load_promot(name):
with open(os.path.join(CUR_DIR, f'promot/{name}.md'), 'r') as f:
return f.read()
def ask(input:str, p_name:str, stream=False):
his = [{"role":"system", "content": load_promot(p_name)}]
his.append({"role":"user", "content": input})
payload = {
"model": MODEL,
"messages": his,
"temperature": 0,
"stream": stream
}
response = requests.post(LLM_URL, headers=HEADERS, json=payload, stream=stream)
if not stream:
return response.json()["choices"][0]["message"]["content"]
else:
# 处理流式响应
full_content = ""
for chunk in response.iter_lines():
if chunk:
# 通常流式响应是SSE格式data: {...}
decoded_chunk = chunk.decode('utf-8')
if decoded_chunk.startswith("data:"):
json_str = decoded_chunk[5:].strip()
if json_str == "[DONE]":
break
try:
chunk_data = json.loads(json_str)
if "choices" in chunk_data and chunk_data["choices"]:
delta = chunk_data["choices"][0].get("delta", {})
if "content" in delta:
print(delta["content"])
full_content += delta["content"]
except json.JSONDecodeError:
continue
return full_content
def work_chain(input:str, t_key:str):
pdict = {"state": "progress", "steps": [{"state":"ok", "msg":"正在生成查询语句"}]}
cache.set(t_key, pdict)
res_text = ask(input, 'w_sql')
if res_text == '请以 查询 开头,重新描述你的需求':
pdict["state"] = "error"
pdict["steps"].append({"state":"error", "msg":res_text})
cache.set(t_key, pdict)
return
else:
pdict["steps"].append({"state":"ok", "msg":"查询语句生成成功", "content":res_text})
cache.set(t_key, pdict)
if not is_safe_sql(res_text):
pdict["state"] = "error"
pdict["steps"].append({"state":"error", "msg":"当前查询存在风险,请重新描述你的需求"})
cache.set(t_key, pdict)
return
pdict["steps"].append({"state":"ok", "msg":"正在执行查询语句"})
cache.set(t_key, pdict)
res = execute_raw_sql(res_text)
pdict["steps"].append({"state":"ok", "msg":"查询语句执行成功", "content":res})
cache.set(t_key, pdict)
pdict["steps"].append({"state":"ok", "msg":"正在生成报告"})
cache.set(t_key, pdict)
res2 = ask(json.dumps(res, cls=MyJSONEncoder, ensure_ascii=False), 'w_ana')
content = res2.lstrip('```html ').rstrip('```')
pdict["state"] = "done"
pdict["content"] = content
pdict["steps"].append({"state":"ok", "msg":"报告生成成功", "content": content})
cache.set(t_key, pdict)
return
class InputSerializer(serializers.Serializer):
input = serializers.CharField(label="查询需求")
class WorkChain(MyLoggingMixin, APIView):
@swagger_auto_schema(
operation_summary="提交查询需求",
request_body=InputSerializer)
def post(self, request):
llm_enabled = getattr(settings, "LLM_ENABLED", False)
if not llm_enabled:
raise ParseError('LLM功能未启用')
input = request.data.get('input')
t_key = f'ichat_{uuid.uuid4()}'
MyThread(target=work_chain, args=(input, t_key)).start()
return Response({'ichat_tid': t_key})
@swagger_auto_schema(
operation_summary="获取查询进度")
def get(self, request):
llm_enabled = getattr(settings, "LLM_ENABLED", False)
if not llm_enabled:
raise ParseError('LLM功能未启用')
ichat_tid = request.GET.get('ichat_tid')
if ichat_tid:
return Response(cache.get(ichat_tid))
if __name__ == "__main__":
print(work_chain("查询 一次超洗 工段在2025年6月的生产合格数等并形成报告"))
from apps.ichat.views2 import work_chain
print(work_chain('查询外观检验工段在2025年6月的生产合格数等并形成报告'))