feat: ichat /增加ichat模块，优化对话流处理
This commit is contained in:
parent
a4ba33550e
commit
f3ab4476a4
|
|
@ -12,6 +12,6 @@ class Message(BaseModel):
|
|||
"""
|
||||
TN: 消息
|
||||
"""
|
||||
conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, verbose_name='对话')
|
||||
conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, related_name='messages', verbose_name='对话')
|
||||
content = models.TextField(verbose_name='消息内容')
|
||||
role = models.CharField("角色", max_length=10, default='user', help_text="system/user")
|
||||
|
|
|
|||
|
|
@ -0,0 +1,22 @@
|
|||
import json
|
||||
from .models import Message
|
||||
from django.http import StreamingHttpResponse
|
||||
|
||||
def stream_generator(stream_response: bytes, conversation_id: str):
|
||||
full_content = ''
|
||||
for chunk in stream_response.iter_content(chunk_size=1024):
|
||||
if chunk:
|
||||
full_content += chunk.decode('utf-8')
|
||||
try:
|
||||
data = json.loads(full_content)
|
||||
content = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
|
||||
Message.objects.create(
|
||||
conversation_id=conversation_id,
|
||||
content=content
|
||||
)
|
||||
yield f" data:{content}\n\n"
|
||||
full_content = ''
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
return StreamingHttpResponse(stream_generator(stream_response, conversation_id), content_type='text/event-stream')
|
||||
|
||||
|
|
@ -6,7 +6,7 @@ from apps.utils.constants import EXCLUDE_FIELDS
|
|||
class MessageSerializer(serializers.ModelSerializer):
|
||||
class Meta:
|
||||
model = Message
|
||||
fields = ['id', 'conversation', 'mode', 'content', 'role']
|
||||
fields = ['id', 'conversation', 'content', 'role']
|
||||
read_only_fields = EXCLUDE_FIELDS
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,10 +1,14 @@
|
|||
|
||||
from django.urls import path
|
||||
from apps.ichat.views import QueryLLMview, ConversationView
|
||||
from django.urls import path, include
|
||||
from rest_framework.routers import DefaultRouter
|
||||
from apps.ichat.views import QueryLLMviewSet, ConversationViewSet
|
||||
|
||||
API_BASE_URL = 'api/ichat/'
|
||||
|
||||
router = DefaultRouter()
|
||||
|
||||
router.register('conversation', ConversationViewSet, basename='conversation')
|
||||
router.register('message', QueryLLMviewSet, basename='message')
|
||||
urlpatterns = [
|
||||
path(API_BASE_URL + 'query/', QueryLLMview.as_view(), name='llm_query'),
|
||||
path(API_BASE_URL + 'conversation/', ConversationView.as_view(), name='conversation')
|
||||
|
||||
]
|
||||
path(API_BASE_URL, include(router.urls)),
|
||||
]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,88 @@
|
|||
import re
|
||||
import psycopg2
|
||||
import threading
|
||||
from django.db import transaction
|
||||
from .models import Message
|
||||
|
||||
# 数据库连接
|
||||
def connect_db():
|
||||
from server.conf import DATABASES
|
||||
db_conf = DATABASES['default']
|
||||
conn = psycopg2.connect(
|
||||
host=db_conf['HOST'],
|
||||
port=db_conf['PORT'],
|
||||
user=db_conf['USER'],
|
||||
password=db_conf['PASSWORD'],
|
||||
database=db_conf['NAME']
|
||||
)
|
||||
return conn
|
||||
|
||||
def extract_sql_code(text):
|
||||
# 优先尝试 ```sql 包裹的语句
|
||||
match = re.search(r"```sql\s*(.+?)```", text, re.DOTALL | re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
|
||||
# fallback: 寻找首个 select 语句
|
||||
match = re.search(r"(SELECT\s.+?;)", text, re.IGNORECASE | re.DOTALL)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_schema_text(conn, table_names:list):
|
||||
cur = conn.cursor()
|
||||
query = """
|
||||
SELECT
|
||||
table_name, column_name, data_type
|
||||
FROM
|
||||
information_schema.columns
|
||||
WHERE
|
||||
table_schema = 'public'
|
||||
and table_name in %s;
|
||||
"""
|
||||
cur.execute(query, (tuple(table_names), ))
|
||||
|
||||
schema = {}
|
||||
for table_name, column_name, data_type in cur.fetchall():
|
||||
if table_name not in schema:
|
||||
schema[table_name] = []
|
||||
schema[table_name].append(f"{column_name} ({data_type})")
|
||||
cur.close()
|
||||
schema_text = ""
|
||||
for table_name, columns in schema.items():
|
||||
schema_text += f"表{table_name} 包含列：{', '.join(columns)}\n"
|
||||
return schema_text
|
||||
|
||||
|
||||
def is_safe_sql(sql:str) -> bool:
|
||||
sql = sql.strip().lower()
|
||||
return sql.startswith("select") or sql.startswith("show") and not re.search(r"delete|update|insert|drop|create|alter", sql)
|
||||
|
||||
def execute_sql(conn, sql_query):
|
||||
cur = conn.cursor()
|
||||
cur.execute(sql_query)
|
||||
try:
|
||||
rows = cur.fetchall()
|
||||
columns = [desc[0] for desc in cur.description]
|
||||
result = [dict(zip(columns, row)) for row in rows]
|
||||
except psycopg2.ProgrammingError:
|
||||
result = cur.statusmessage
|
||||
cur.close()
|
||||
return result
|
||||
|
||||
def strip_sql_markdown(content: str) -> str:
|
||||
# 去掉包裹在 ```sql 或 ``` 中的内容
|
||||
match = re.search(r"```sql\s*(.*?)```", content, re.DOTALL | re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
else:
|
||||
return None
|
||||
|
||||
# ORM 写入包装函数
|
||||
def save_message_thread_safe(**kwargs):
|
||||
def _save():
|
||||
with transaction.atomic():
|
||||
Message.objects.create(**kwargs)
|
||||
threading.Thread(target=_save).start()
|
||||
|
|
@ -1,173 +1,155 @@
|
|||
import requests
|
||||
import psycopg2
|
||||
import json
|
||||
from rest_framework.views import APIView
|
||||
from apps.ichat.serializers import MessageSerializer, ConversationSerializer
|
||||
from rest_framework.response import Response
|
||||
from ichat.models import Conversation, Message
|
||||
from rest_framework.generics import get_object_or_404
|
||||
#本地部署模型
|
||||
from apps.ichat.models import Conversation, Message
|
||||
from apps.ichat.utils import connect_db, extract_sql_code, execute_sql, get_schema_text, is_safe_sql, save_message_thread_safe
|
||||
from django.http import StreamingHttpResponse, JsonResponse
|
||||
from rest_framework.decorators import action
|
||||
from apps.utils.viewsets import CustomGenericViewSet, CustomModelViewSet
|
||||
|
||||
# API_KEY = "sk-5644e2d6077b46b9a04a8a2b12d6b693"
|
||||
# API_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
# MODEL = "qwen-plus"
|
||||
|
||||
#本地部署的模式
|
||||
# API_KEY = "JJVAide0hw3eaugGmxecyYYFw45FX2LfhnYJtC+W2rw"
|
||||
# API_BASE = "http://106.0.4.200:9000/v1"
|
||||
# MODEL = "Qwen/Qwen2.5-14B-Instruct"
|
||||
API_KEY = "JJVAide0hw3eaugGmxecyYYFw45FX2LfhnYJtC+W2rw"
|
||||
API_BASE = "http://106.0.4.200:9000/v1"
|
||||
MODEL = "qwen14b"
|
||||
|
||||
# google gemini
|
||||
API_KEY = "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"
|
||||
API_BASE = "https://openrouter.ai/api/v1"
|
||||
MODEL="google/gemini-2.0-flash-exp:free"
|
||||
# API_KEY = "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"
|
||||
# API_BASE = "https://openrouter.ai/api/v1"
|
||||
# MODEL="google/gemini-2.0-flash-exp:free"
|
||||
|
||||
# deepseek v3
|
||||
# API_KEY = "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"
|
||||
# API_BASE = "https://openrouter.ai/api/v1"
|
||||
# MODEL="deepseek/deepseek-chat-v3-0324:free"
|
||||
|
||||
|
||||
TABLES = ["enm_mpoint", "enm_mpointstat", "enm_mplogx"] # 如果整个数据库全都给模型，准确率下降，所以只给模型部分表
|
||||
# 数据库连接
|
||||
def connect_db():
|
||||
from server.conf import DATABASES
|
||||
db_conf = DATABASES['default']
|
||||
conn = psycopg2.connect(
|
||||
host=db_conf['HOST'],
|
||||
port=db_conf['PORT'],
|
||||
user=db_conf['USER'],
|
||||
password=db_conf['PASSWORD'],
|
||||
database=db_conf['NAME']
|
||||
)
|
||||
return conn
|
||||
|
||||
def get_schema_text(conn, table_names:list):
|
||||
cur = conn.cursor()
|
||||
query = """
|
||||
SELECT
|
||||
table_name, column_name, data_type
|
||||
FROM
|
||||
information_schema.columns
|
||||
WHERE
|
||||
table_schema = 'public'
|
||||
and table_name in %s;
|
||||
"""
|
||||
cur.execute(query, (tuple(table_names), ))
|
||||
|
||||
schema = {}
|
||||
for table_name, column_name, data_type in cur.fetchall():
|
||||
if table_name not in schema:
|
||||
schema[table_name] = []
|
||||
schema[table_name].append(f"{column_name} ({data_type})")
|
||||
cur.close()
|
||||
schema_text = ""
|
||||
for table_name, columns in schema.items():
|
||||
schema_text += f"表{table_name} 包含列：{', '.join(columns)}\n"
|
||||
return schema_text
|
||||
|
||||
|
||||
# 调用大模型生成sql
|
||||
def call_llm_api(prompt, api_key=API_KEY, api_base=API_BASE, model=MODEL):
|
||||
url = f"{api_base}/chat/completions"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}"
|
||||
}
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": 0,
|
||||
}
|
||||
response = requests.post(url, headers=headers, json=payload)
|
||||
response.raise_for_status()
|
||||
print("\n大模型返回：\n", response.json())
|
||||
return response.json()["choices"][0]["message"]["content"]
|
||||
class QueryLLMviewSet(CustomModelViewSet):
|
||||
queryset = Message.objects.all()
|
||||
serializer_class = MessageSerializer
|
||||
ordering = ['create_time']
|
||||
perms_map = {'get':'*', 'post':'*', 'put':'*'}
|
||||
|
||||
|
||||
def execute_sql(conn, sql_query):
|
||||
cur = conn.cursor()
|
||||
cur.execute(sql_query)
|
||||
try:
|
||||
rows = cur.fetchall()
|
||||
columns = [desc[0] for desc in cur.description]
|
||||
result = [dict(zip(columns, row)) for row in rows]
|
||||
except psycopg2.ProgrammingError:
|
||||
result = cur.statusmessage
|
||||
cur.close()
|
||||
return result
|
||||
|
||||
|
||||
def strip_sql_markdown(content: str) -> str:
|
||||
import re
|
||||
# 去掉包裹在 ```sql 或 ``` 中的内容
|
||||
match = re.search(r"```sql\s*(.*?)```", content, re.DOTALL | re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
class QueryLLMview(APIView):
|
||||
def post(self, request):
|
||||
serializer = MessageSerializer(data=request.data)
|
||||
@action(methods=['post'], detail=False, perms_map={'post':'*'} ,serializer_class=MessageSerializer)
|
||||
def completion(self, request):
|
||||
serializer = self.get_serializer(data=request.data)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
serializer.save()
|
||||
prompt = serializer.validated_data['prompt']
|
||||
conn = connect_db()
|
||||
# 数据库表结构
|
||||
schema_text = get_schema_text(conn, TABLES)
|
||||
user_prompt = f"""你是可能是一个专业的数据库工程师，根据以下数据库结构：
|
||||
prompt = serializer.validated_data['content']
|
||||
conversation = serializer.validated_data['conversation']
|
||||
if not prompt or not conversation:
|
||||
return JsonResponse({"error": "缺少 prompt 或 conversation"}, status=400)
|
||||
save_message_thread_safe(content=prompt, conversation=conversation, role="user")
|
||||
url = f"{API_BASE}/chat/completions"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {API_KEY}"
|
||||
}
|
||||
|
||||
user_prompt = f"""
|
||||
请判断以下问题是否与数据库查询或操作相关。如果是，回答"database"；如果不是，回答"general"。
|
||||
|
||||
问题: {prompt}
|
||||
|
||||
只需回答"database"或"general"，不要有其他内容。
|
||||
"""
|
||||
_payload = {
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": user_prompt}],
|
||||
"temperature": 0,
|
||||
"max_tokens": 10
|
||||
}
|
||||
try:
|
||||
class_response = requests.post(url, headers=headers, json=_payload)
|
||||
class_response.raise_for_status()
|
||||
class_result = class_response.json()
|
||||
question_type = class_result.get('choices', [{}])[0].get('message', {}).get('content', '').strip().lower()
|
||||
if question_type == "database":
|
||||
conn = connect_db()
|
||||
schema_text = get_schema_text(conn, TABLES)
|
||||
user_prompt = f"""你是一个专业的数据库工程师，根据以下数据库结构：
|
||||
{schema_text}
|
||||
请根据我的需求生成一条标准的PostgreSQL SQL语句，直接返回SQL，不要额外解释。
|
||||
需求是：{prompt}
|
||||
"""
|
||||
llm_data = call_llm_api(user_prompt)
|
||||
# 判断是否生成的是sql 如果不是直接返回message
|
||||
generated_sql = strip_sql_markdown(llm_data)
|
||||
if generated_sql:
|
||||
try:
|
||||
result = execute_sql(conn, generated_sql)
|
||||
return Response({"result": result})
|
||||
except Exception as e:
|
||||
print("\n第一次执行SQL报错了，错误信息：", str(e))
|
||||
# 如果第一次执行SQL报错，则重新生成SQL
|
||||
fix_prompt = f"""刚才你生成的SQL出现了错误，错误信息是：{str(e)}
|
||||
请根据这个错误修正你的SQL，返回新的正确的SQL，直接给出SQL，不要解释。
|
||||
数据库结构如下：
|
||||
{schema_text}
|
||||
用户需求是：{prompt}
|
||||
"""
|
||||
fixed_sql = call_llm_api(fix_prompt)
|
||||
fixed_sql = strip_sql_markdown(fixed_sql)
|
||||
try:
|
||||
results = execute_sql(conn, fixed_sql)
|
||||
print("\n修正后的查询结果：")
|
||||
print(results)
|
||||
return Response({"result": results})
|
||||
except Exception as e2:
|
||||
print("\n修正后的SQL仍然报错，错误信息：", str(e2))
|
||||
return Response({"error": "SQL执行失败", "detail": str(e2)}, status=400)
|
||||
finally:
|
||||
conn.close()
|
||||
else:
|
||||
return Response({"result": llm_data})
|
||||
"""
|
||||
else:
|
||||
user_prompt = f"""
|
||||
回答以下问题，不需要涉及数据库查询：
|
||||
|
||||
问题: {prompt}
|
||||
|
||||
请直接回答问题，不要提及数据库或SQL。
|
||||
"""
|
||||
# TODO 是否应该拿到conservastion的id，然后根据id去数据库查询所以的messages, 然后赋值给messages
|
||||
history = Message.objects.filter(conversation=conversation).order_by('create_time')
|
||||
chat_history = [{"role": msg.role, "content": msg.content} for msg in history]
|
||||
chat_history.append({"role": "user", "content": prompt})
|
||||
print("chat_history", chat_history)
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"messages": chat_history,
|
||||
"temperature": 0,
|
||||
"stream": True
|
||||
}
|
||||
response = requests.post(url, headers=headers, json=payload)
|
||||
response.raise_for_status()
|
||||
except requests.exceptions.RequestException as e:
|
||||
return JsonResponse({"error":f"LLM API调用失败: {e}"}, status=500)
|
||||
def stream_generator():
|
||||
accumulated_content = ""
|
||||
for line in response.iter_lines():
|
||||
if line:
|
||||
decoded_line = line.decode('utf-8')
|
||||
if decoded_line.startswith('data:'):
|
||||
if decoded_line.strip() == "data: [DONE]":
|
||||
break # OpenAI-style标志结束
|
||||
try:
|
||||
data = json.loads(decoded_line[6:])
|
||||
content = data.get('choices', [{}])[0].get('delta', {}).get('content', '')
|
||||
print("content", content)
|
||||
if content:
|
||||
accumulated_content += content
|
||||
yield f"data: {content}\n\n"
|
||||
|
||||
except Exception as e:
|
||||
yield f"data: [解析失败]: {str(e)}\n\n"
|
||||
print("accumulated_content", accumulated_content)
|
||||
print("question_type", question_type)
|
||||
print("conversation", conversation)
|
||||
save_message_thread_safe(content=accumulated_content, conversation=conversation, role="system")
|
||||
|
||||
if question_type == "database":
|
||||
sql = extract_sql_code(accumulated_content)
|
||||
if sql:
|
||||
try:
|
||||
conn = connect_db()
|
||||
if is_safe_sql(sql):
|
||||
result = execute_sql(conn, sql)
|
||||
save_message_thread_safe(content=f"SQL结果: {result}", conversation=conversation, role="system")
|
||||
yield f"data: SQL执行结果: {result}\n\n"
|
||||
else:
|
||||
yield f"data: 拒绝执行非查询类 SQL：{sql}\n\n"
|
||||
except Exception as e:
|
||||
yield f"data: SQL执行失败: {str(e)}\n\n"
|
||||
finally:
|
||||
if conn:
|
||||
conn.close()
|
||||
else:
|
||||
yield "data: \\n[文本结束]\n\n"
|
||||
return StreamingHttpResponse(stream_generator(), content_type='text/event-stream')
|
||||
|
||||
|
||||
# 先新建对话 生成对话session_id
|
||||
class ConversationView(APIView):
|
||||
def get(self, request):
|
||||
conversation = Conversation.objects.all()
|
||||
serializer = ConversationSerializer(conversation, many=True)
|
||||
return Response(serializer.data)
|
||||
class ConversationViewSet(CustomModelViewSet):
|
||||
queryset = Conversation.objects.all()
|
||||
serializer_class = ConversationSerializer
|
||||
ordering = ['create_time']
|
||||
perms_map = {'get':'*', 'post':'*', 'put':'*'}
|
||||
|
||||
def post(self, request):
|
||||
serializer = ConversationSerializer(data=request.data)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
serializer.save()
|
||||
return Response(serializer.data)
|
||||
|
||||
def put(self, request, pk):
|
||||
conversation = get_object_or_404(Conversation, pk=pk)
|
||||
serializer = ConversationSerializer(conversation, data=request.data)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
serializer.save()
|
||||
return Response(serializer.data)
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
|
||||
from django.core.cache import cache
|
||||
from django.http import StreamingHttpResponse
|
||||
from rest_framework.decorators import action
|
||||
from rest_framework.exceptions import ParseError
|
||||
from rest_framework.mixins import RetrieveModelMixin
|
||||
|
|
@ -58,6 +59,9 @@ class CustomGenericViewSet(MyLoggingMixin, GenericViewSet):
|
|||
return super().__new__(cls)
|
||||
|
||||
def finalize_response(self, request, response, *args, **kwargs):
|
||||
# 如果是流式响应，直接返回
|
||||
if isinstance(response, StreamingHttpResponse):
|
||||
return response
|
||||
if self.hash_k and self.cache_seconds:
|
||||
cache.set(self.hash_k, response.data,
|
||||
timeout=self.cache_seconds) # 将结果存入缓存，设置超时时间
|
||||
|
|
|
|||
|
|
@ -63,6 +63,7 @@ INSTALLED_APPS = [
|
|||
'apps.wf',
|
||||
'apps.ecm',
|
||||
'apps.hrm',
|
||||
'apps.ichat',
|
||||
'apps.am',
|
||||
'apps.vm',
|
||||
'apps.rpm',
|
||||
|
|
|
|||
|
|
@ -44,7 +44,7 @@ urlpatterns = [
|
|||
|
||||
# api
|
||||
path('', include('apps.auth1.urls')),
|
||||
# path('', include('apps.ichat.urls')),
|
||||
path('', include('apps.ichat.urls')),
|
||||
path('', include('apps.system.urls')),
|
||||
path('', include('apps.monitor.urls')),
|
||||
path('', include('apps.wf.urls')),
|
||||
|
|
|
|||
Loading…
Reference in New Issue