daren/apps/chat/consumers.py

# apps/chat/consumers.py
from channels.generic.websocket import AsyncWebsocketConsumer
import json
from channels.db import database_sync_to_async
from apps.chat.models import ChatHistory
from apps.user.models import UserToken
from django.conf import settings
import logging
import traceback
import uuid
import aiohttp
from urllib.parse import parse_qs
from django.utils import timezone
from rest_framework.permissions import IsAuthenticated
from apps.user.authentication import CustomTokenAuthentication
logger = logging.getLogger(__name__)

class ChatStreamConsumer(AsyncWebsocketConsumer):

    # 固定知识库ID
    DEFAULT_KNOWLEDGE_BASE_ID = "b680a4fa-37be-11f0-a7cb-0242ac120002"
    
    async def connect(self):
        """建立WebSocket连接"""
        try:
            # 从URL参数中获取token
            query_string = self.scope.get('query_string', b'').decode()
            query_params = parse_qs(query_string)
            token_key = query_params.get('token', [''])[0]
            
            if not token_key:
                logger.warning("WebSocket连接尝试，但没有提供token")
                await self.close()
                return
                
            # 验证token
            self.user = await self.get_user_from_token(token_key)
            if not self.user:
                logger.warning(f"WebSocket连接尝试，但token无效: {token_key}")
                await self.close()
                return

            # 将用户信息存储在scope中
            self.scope["user"] = self.user
            await self.accept()
            logger.info(f"用户 {self.user.email} 流式输出WebSocket连接成功")
            
        except Exception as e:
            logger.error(f"WebSocket连接错误: {str(e)}")
            await self.close()

    @database_sync_to_async
    def get_user_from_token(self, token_key):
        try:
            # 使用项目的UserToken模型而不是rest_framework的Token
            token = UserToken.objects.select_related('user').get(
                token=token_key,
                expired_at__gt=timezone.now()  # 确保token未过期
            )
            return token.user
        except UserToken.DoesNotExist:
            return None

    async def disconnect(self, close_code):
        """关闭WebSocket连接"""
        logger.info(f"用户 {self.user.email if hasattr(self, 'user') else 'unknown'} WebSocket连接断开，代码: {close_code}")

    async def receive(self, text_data):
        """接收消息并处理"""
        try:
            data = json.loads(text_data)
            
            # 检查必填字段
            if 'question' not in data:
                await self.send_error("缺少必填字段: question")
                return
                
            if 'conversation_id' not in data:
                await self.send_error("缺少必填字段: conversation_id")
                return

            # 处理新会话或现有会话
            await self.process_chat_request(data)
                
        except Exception as e:
            logger.error(f"处理消息时出错: {str(e)}")
            logger.error(traceback.format_exc())
            await self.send_error(f"处理消息时出错: {str(e)}")

    async def process_chat_request(self, data):
        """处理聊天请求"""
        try:
            conversation_id = data['conversation_id']
            question = data['question']
            
            # 准备metadata
            metadata = {}
            
            # 创建问题记录
            question_record = await self.create_question_record(
                conversation_id, 
                question, 
                metadata
            )
            
            if not question_record:
                return
                
            # 创建AI回答记录
            answer_record = await self.create_answer_record(
                conversation_id,
                question_record,
                metadata
            )
            
            # 发送初始响应
            await self.send_json({
                'code': 200,
                'message': '开始流式传输',
                'data': {
                    'id': str(answer_record.id),
                    'conversation_id': str(conversation_id),
                    'content': '',
                    'is_end': False
                }
            })
            
            # 设置外部API需要的ID列表 - 简化为空列表
            dataset_external_id_list = []
            
            # 调用外部API获取流式响应
            await self.stream_from_external_api(
                conversation_id,
                question,
                dataset_external_id_list,
                answer_record,
                metadata
            )
            
        except Exception as e:
            logger.error(f"处理聊天请求时出错: {str(e)}")
            logger.error(traceback.format_exc())
            await self.send_error(f"处理聊天请求时出错: {str(e)}")

    @database_sync_to_async
    def create_question_record(self, conversation_id, question, metadata):
        """创建问题记录"""
        try:
            title = "New chat"
            
            # 创建用户问题记录
            return ChatHistory.objects.create(
                user=self.scope["user"],
                knowledge_base_id=self.DEFAULT_KNOWLEDGE_BASE_ID,
                conversation_id=str(conversation_id),
                title=title,
                role='user',
                content=question,
                metadata=metadata
            )
        except Exception as e:
            logger.error(f"创建问题记录时出错: {str(e)}")
            return None

    @database_sync_to_async
    def create_answer_record(self, conversation_id, question_record, metadata):
        """创建AI回答记录"""
        try:
            return ChatHistory.objects.create(
                user=self.scope["user"],
                knowledge_base_id=self.DEFAULT_KNOWLEDGE_BASE_ID,
                conversation_id=str(conversation_id),
                title=question_record.title,
                parent_id=str(question_record.id),
                role='assistant',
                content="",  # 初始内容为空
                metadata=metadata
            )
        except Exception as e:
            logger.error(f"创建回答记录时出错: {str(e)}")
            return None

    async def stream_from_external_api(self, conversation_id, question, dataset_external_id_list, answer_record, metadata):
        """从外部API获取流式响应"""
        try:
            # 获取标题
            title = answer_record.title or 'New chat'
            
            # 异步收集完整内容，用于最后保存
            full_content = ""
            
            # 使用aiohttp进行异步HTTP请求
            async with aiohttp.ClientSession() as session:
                # 第一步: 创建聊天会话
                async with session.post(
                    f"{settings.API_BASE_URL}/api/application/chat/open",
                    json={
                        "id": "d5d11efa-ea9a-11ef-9933-0242ac120006",
                        "model_id": metadata.get('model_id', '7a214d0e-e65e-11ef-9f4a-0242ac120006'),
                        "dataset_id_list": dataset_external_id_list,
                        "multiple_rounds_dialogue": False,
                        "dataset_setting": {
                            "top_n": 10, "similarity": "0.3",
                            "max_paragraph_char_number": 10000,
                            "search_mode": "blend",
                            "no_references_setting": {
                                "value": "{question}",
                                "status": "ai_questioning"
                            }
                        },
                        "model_setting": {
                            "prompt": "**相关文档内容**：{data} **回答要求**：如果相关文档内容中没有可用信息，请回答\"没有在知识库中查找到相关信息，建议咨询相关技术支持或参考官方文档进行操作\"。请根据相关文档内容回答用户问题。不要输出与用户问题无关的内容。请使用中文回答客户问题。**用户问题**：{question}"
                        },
                        "problem_optimization": False
                    }
                ) as chat_response:
                    
                    if chat_response.status != 200:
                        error_msg = f"外部API调用失败: {await chat_response.text()}"
                        logger.error(error_msg)
                        await self.send_error(error_msg)
                        return
                    
                    chat_data = await chat_response.json()
                    if chat_data.get('code') != 200 or not chat_data.get('data'):
                        error_msg = f"外部API返回错误: {chat_data}"
                        logger.error(error_msg)
                        await self.send_error(error_msg)
                        return
                    
                    chat_id = chat_data['data']
                    logger.info(f"成功创建聊天会话, chat_id: {chat_id}")
                    
                    # 第二步: 建立流式连接
                    message_url = f"{settings.API_BASE_URL}/api/application/chat_message/{chat_id}"
                    logger.info(f"开始流式请求: {message_url}")
                    
                    # 创建流式请求
                    async with session.post(
                        url=message_url,
                        json={"message": question, "re_chat": False, "stream": True},
                        headers={"Content-Type": "application/json"}
                    ) as message_request:
                        
                        if message_request.status != 200:
                            error_msg = f"外部API聊天消息调用失败: {message_request.status}, {await message_request.text()}"
                            logger.error(error_msg)
                            await self.send_error(error_msg)
                            return
                        
                        # 创建一个缓冲区以处理分段的数据
                        buffer = ""
                        
                        # 读取并处理每个响应块
                        logger.info("开始处理流式响应")
                        async for chunk in message_request.content.iter_any():
                            chunk_str = chunk.decode('utf-8')
                            buffer += chunk_str
                            
                            # 检查是否有完整的数据行
                            while '\n\n' in buffer:
                                parts = buffer.split('\n\n', 1)
                                line = parts[0]
                                buffer = parts[1]
                                
                                if line.startswith('data: '):
                                    try:
                                        # 提取JSON数据
                                        json_str = line[6:]  # 去掉 "data: " 前缀
                                        data = json.loads(json_str)
                                        
                                        # 记录并处理部分响应
                                        if 'content' in data:
                                            content_part = data['content']
                                            full_content += content_part
                                            
                                            # 发送部分内容
                                            await self.send_json({
                                                'code': 200,
                                                'message': 'partial',
                                                'data': {
                                                    'id': str(answer_record.id),
                                                    'conversation_id': str(conversation_id),
                                                    'content': content_part,
                                                    'is_end': data.get('is_end', False)
                                                }
                                            })
                                            
                                        # 处理结束标记
                                        if data.get('is_end', False):
                                            logger.info("收到流式响应结束标记")
                                            # 保存完整内容
                                            await self.update_answer_content(answer_record.id, full_content.strip())
                                            
                                            # 处理标题
                                            title = await self.get_or_generate_title(
                                                conversation_id,
                                                question,
                                                full_content.strip()
                                            )
                                            
                                            # 发送最终响应
                                            await self.send_json({
                                                'code': 200,
                                                'message': '完成',
                                                'data': {
                                                    'id': str(answer_record.id),
                                                    'conversation_id': str(conversation_id),
                                                    'title': title,
                                                    'role': 'assistant',
                                                    'content': full_content.strip(),
                                                    'created_at': answer_record.created_at.strftime('%Y-%m-%d %H:%M:%S'),
                                                    'is_end': True
                                                }
                                            })
                                            return
                                            
                                    except json.JSONDecodeError as e:
                                        logger.error(f"JSON解析错误: {e}, 数据: {line}")
                                        continue
        
        except Exception as e:
            logger.error(f"流式处理出错: {str(e)}")
            logger.error(traceback.format_exc())
            await self.send_error(str(e))
            
            # 保存已收集的内容
            if 'full_content' in locals() and full_content:
                try:
                    await self.update_answer_content(answer_record.id, full_content.strip())
                except Exception as save_error:
                    logger.error(f"保存部分内容失败: {str(save_error)}")

    @database_sync_to_async
    def update_answer_content(self, answer_id, content):
        """更新回答内容"""
        try:
            answer_record = ChatHistory.objects.get(id=answer_id)
            answer_record.content = content
            answer_record.save()
            return True
        except Exception as e:
            logger.error(f"更新回答内容失败: {str(e)}")
            return False

    @database_sync_to_async
    def get_or_generate_title(self, conversation_id, question, answer):
        """获取或生成对话标题"""
        try:
            # 先检查是否已有标题
            current_title = ChatHistory.objects.filter(
                conversation_id=str(conversation_id)
            ).exclude(
                title__in=["New chat", "新对话", ""]
            ).values_list('title', flat=True).first()
            
            if current_title:
                return current_title
                
            # 简单的标题生成逻辑 (可替换为调用DeepSeek API生成标题)
            generated_title = question[:20] + "..." if len(question) > 20 else question
            
            # 更新所有相关记录的标题
            ChatHistory.objects.filter(
                conversation_id=str(conversation_id)
            ).update(title=generated_title)
            
            return generated_title
            
        except Exception as e:
            logger.error(f"获取或生成标题失败: {str(e)}")
            return "新对话"

    async def send_json(self, content):
        """发送JSON格式的消息"""
        await self.send(text_data=json.dumps(content))

    async def send_error(self, message):
        """发送错误消息"""
        await self.send_json({
            'code': 500,
            'message': message,
            'data': {'is_end': True}
        })