feat: ichat /增加ichat模块，优化对话流处理

apps/ichat/models.py

@@ -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")

apps/ichat/script.py

@@ -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')
+    

apps/ichat/serializers.py

@@ -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
 
 

apps/ichat/urls.py

@@ -1,10 +1,14 @@
 
-from django.urls import path
+from django.urls import path, include
-from apps.ichat.views import QueryLLMview, ConversationView
+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, include(router.urls)),
-    path(API_BASE_URL + 'conversation/', ConversationView.as_view(), name='conversation')
+]
-    
-]

apps/ichat/utils.py

@@ -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()

apps/ichat/views.py

@@ -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 apps.ichat.models import Conversation, Message
-from rest_framework.generics import get_object_or_404
+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_KEY = "JJVAide0hw3eaugGmxecyYYFw45FX2LfhnYJtC+W2rw"
-# API_BASE = "http://106.0.4.200:9000/v1"
+API_BASE = "http://106.0.4.200:9000/v1"
-# MODEL = "Qwen/Qwen2.5-14B-Instruct"
+MODEL = "qwen14b"
 
 # google gemini
-API_KEY = "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"
+# API_KEY = "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"
-API_BASE = "https://openrouter.ai/api/v1"
+# API_BASE = "https://openrouter.ai/api/v1"
-MODEL="google/gemini-2.0-flash-exp:free"
+# 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
+class QueryLLMviewSet(CustomModelViewSet):
-def call_llm_api(prompt, api_key=API_KEY, api_base=API_BASE, model=MODEL):
+    queryset = Message.objects.all()
-    url = f"{api_base}/chat/completions"
+    serializer_class = MessageSerializer
-    headers = {
+    ordering = ['create_time']
-        "Content-Type": "application/json",
+    perms_map = {'get':'*', 'post':'*', 'put':'*'}
-        "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"]
 
-
+    @action(methods=['post'], detail=False, perms_map={'post':'*'} ,serializer_class=MessageSerializer)
-def execute_sql(conn, sql_query):
+    def completion(self, request):
-    cur = conn.cursor()
+        serializer = self.get_serializer(data=request.data)
-    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)
         serializer.is_valid(raise_exception=True)
         serializer.save()
-        prompt = serializer.validated_data['prompt']
+        prompt = serializer.validated_data['content']
-        conn = connect_db()
+        conversation = serializer.validated_data['conversation']
-        # 数据库表结构
+        if not prompt or not conversation:
-        schema_text = get_schema_text(conn, TABLES)
+            return JsonResponse({"error": "缺少 prompt 或 conversation"}, status=400)
-        user_prompt = f"""你是可能是一个专业的数据库工程师，根据以下数据库结构：
+        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)
+            else:
-        # 判断是否生成的是sql 如果不是直接返回message
+                user_prompt = f"""
-        generated_sql = strip_sql_markdown(llm_data)
+回答以下问题，不需要涉及数据库查询：
-        if generated_sql:
+
-            try:
+问题: {prompt}
-                result = execute_sql(conn, generated_sql)
+
-                return Response({"result": result})
+请直接回答问题，不要提及数据库或SQL。
-            except Exception as e:
+"""         
-                print("\n第一次执行SQL报错了，错误信息：", str(e))
+            # TODO 是否应该拿到conservastion的id，然后根据id去数据库查询所以的messages, 然后赋值给messages
-                # 如果第一次执行SQL报错，则重新生成SQL
+            history = Message.objects.filter(conversation=conversation).order_by('create_time')
-                fix_prompt = f"""刚才你生成的SQL出现了错误，错误信息是：{str(e)}
+            chat_history = [{"role": msg.role, "content": msg.content} for msg in history]
-    请根据这个错误修正你的SQL，返回新的正确的SQL，直接给出SQL，不要解释。
+            chat_history.append({"role": "user", "content": prompt})
-    数据库结构如下：
+            print("chat_history", chat_history)
-    {schema_text}
+            payload = {
-    用户需求是：{prompt}
+                    "model": MODEL,
-    """
+                    "messages": chat_history,
-                fixed_sql = call_llm_api(fix_prompt)
+                    "temperature": 0,
-                fixed_sql = strip_sql_markdown(fixed_sql)
+                    "stream": True
-                try:
+                }
-                    results = execute_sql(conn, fixed_sql)
+            response = requests.post(url, headers=headers, json=payload)
-                    print("\n修正后的查询结果：")
+            response.raise_for_status()
-                    print(results)
+        except requests.exceptions.RequestException as e:
-                    return Response({"result": results})
+            return JsonResponse({"error":f"LLM API调用失败: {e}"}, status=500)
-                except Exception as e2:
+        def stream_generator():
-                    print("\n修正后的SQL仍然报错，错误信息：", str(e2))
+            accumulated_content = ""
-                    return Response({"error": "SQL执行失败", "detail": str(e2)}, status=400)
+            for line in response.iter_lines():
-            finally:
+                if line:
-                conn.close()
+                    decoded_line = line.decode('utf-8')
-        else:
+                    if decoded_line.startswith('data:'):
-            return Response({"result": llm_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):
+class ConversationViewSet(CustomModelViewSet):
-    def get(self, request):
+    queryset = Conversation.objects.all()
-        conversation = Conversation.objects.all()
+    serializer_class = ConversationSerializer
-        serializer = ConversationSerializer(conversation, many=True)
+    ordering = ['create_time']
-        return Response(serializer.data)
+    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)

apps/utils/viewsets.py

@@ -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)  # 将结果存入缓存，设置超时时间

server/settings.py

@@ -63,6 +63,7 @@ INSTALLED_APPS = [
     'apps.wf',
     'apps.ecm',
     'apps.hrm',
+    'apps.ichat',
     'apps.am',
     'apps.vm',
     'apps.rpm',

server/urls.py

@@ -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')),