import logging import json import traceback import uuid from datetime import datetime from django.db.models import Q, Max, Count from django.http import HttpResponse, StreamingHttpResponse from rest_framework import viewsets, status from rest_framework.permissions import IsAuthenticated from rest_framework.response import Response from rest_framework.decorators import action from apps.accounts.models import User from apps.knowledge_base.models import KnowledgeBase from apps.chat.models import ChatHistory from apps.chat.serializers import ChatHistorySerializer from apps.common.services.chat_service import ChatService from apps.chat.services.chat_api import ( ExternalAPIError, stream_chat_answer, get_chat_answer, generate_conversation_title, get_hit_test_documents, generate_conversation_title_from_deepseek ) from apps.permissions.services.permission_service import KnowledgeBasePermissionMixin logger = logging.getLogger(__name__) class ChatHistoryViewSet(KnowledgeBasePermissionMixin, viewsets.ModelViewSet): permission_classes = [IsAuthenticated] serializer_class = ChatHistorySerializer queryset = ChatHistory.objects.all() def get_queryset(self): """确保用户只能看到自己的未删除的聊天记录以及有权限的知识库关联的聊天记录""" user = self.request.user accessible_kb_ids = [ kb.id for kb in KnowledgeBase.objects.all() if self.check_knowledge_base_permission(kb, user, 'read') ] return ChatHistory.objects.filter( Q(user=user) | Q(knowledge_base_id__in=accessible_kb_ids), is_deleted=False ) def list(self, request): """获取对话列表概览""" try: page = int(request.query_params.get('page', 1)) page_size = int(request.query_params.get('page_size', 10)) latest_chats = self.get_queryset().values( 'conversation_id' ).annotate( latest_id=Max('id'), message_count=Count('id'), last_message=Max('created_at') ).order_by('-last_message') total = latest_chats.count() start = (page - 1) * page_size end = start + page_size chats = latest_chats[start:end] results = [] for chat in chats: latest_record = ChatHistory.objects.get(id=chat['latest_id']) dataset_info = [] if latest_record.metadata: dataset_id_list = latest_record.metadata.get('dataset_id_list', []) dataset_names = latest_record.metadata.get('dataset_names', []) if dataset_id_list: if dataset_names and len(dataset_names) == len(dataset_id_list): dataset_info = [ {'id': str(id), 'name': name} for id, name in zip(dataset_id_list, dataset_names) ] else: datasets = KnowledgeBase.objects.filter(id__in=dataset_id_list) dataset_info = [ {'id': str(ds.id), 'name': ds.name} for ds in datasets ] results.append({ 'conversation_id': chat['conversation_id'], 'message_count': chat['message_count'], 'last_message': latest_record.content, 'last_time': chat['last_message'].strftime('%Y-%m-%d %H:%M:%S'), 'dataset_id_list': [ds['id'] for ds in dataset_info], 'datasets': dataset_info }) return Response({ 'code': 200, 'message': '获取成功', 'data': { 'total': total, 'page': page, 'page_size': page_size, 'results': results } }) except Exception as e: logger.error(f"获取聊天记录失败: {str(e)}") import traceback logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f'获取聊天记录失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @action(detail=False, methods=['get']) def conversation_detail(self, request): """获取特定对话的详细信息""" try: conversation_id = request.query_params.get('conversation_id') if not conversation_id: return Response({ 'code': 400, 'message': '缺少conversation_id参数', 'data': None }, status=status.HTTP_400_BAD_REQUEST) chat_service = ChatService() result = chat_service.get_conversation_detail(request.user, conversation_id) return Response({ 'code': 200, 'message': '获取成功', 'data': result }) except ValueError as e: return Response({ 'code': 404, 'message': str(e), 'data': None }, status=status.HTTP_404_NOT_FOUND) except Exception as e: logger.error(f"获取对话详情失败: {str(e)}") import traceback logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f'获取对话详情失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @action(detail=False, methods=['get']) def available_datasets(self, request): """获取用户可访问的知识库列表""" try: user = request.user accessible_datasets = [ dataset for dataset in KnowledgeBase.objects.all() if self.check_knowledge_base_permission(dataset, user, 'read') ] return Response({ 'code': 200, 'message': '获取成功', 'data': [ { 'id': str(ds.id), 'name': ds.name, 'type': ds.type, 'department': ds.department, 'description': ds.desc } for ds in accessible_datasets ] }) except Exception as e: logger.error(f"获取可用知识库列表失败: {str(e)}") return Response({ 'code': 500, 'message': f'获取可用知识库列表失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @action(detail=False, methods=['post']) def create_conversation(self, request): """创建会话 - 先选择知识库创建会话ID，不发送问题""" try: data = request.data # 检查知识库ID：支持dataset_id或dataset_id_list格式 dataset_ids = [] if 'dataset_id' in data: dataset_id = data['dataset_id'] # 直接使用标准UUID格式 dataset_ids.append(str(dataset_id)) elif 'dataset_id_list' in data and isinstance(data['dataset_id_list'], (list, str)): # 处理可能的字符串格式 if isinstance(data['dataset_id_list'], str): try: # 尝试解析JSON字符串 dataset_list = json.loads(data['dataset_id_list']) if isinstance(dataset_list, list): dataset_ids = [str(id) for id in dataset_list] except json.JSONDecodeError: # 如果解析失败，可能是单个ID dataset_ids = [str(data['dataset_id_list'])] else: # 如果已经是列表，直接使用标准UUID格式 dataset_ids = [str(id) for id in data['dataset_id_list']] else: return Response({ 'code': 400, 'message': '缺少必填字段: dataset_id 或 dataset_id_list', 'data': None }, status=status.HTTP_400_BAD_REQUEST) if not dataset_ids: return Response({ 'code': 400, 'message': '至少需要提供一个知识库ID', 'data': None }, status=status.HTTP_400_BAD_REQUEST) # 验证所有知识库 user = request.user knowledge_bases = [] # 存储所有知识库对象 for kb_id in dataset_ids: try: knowledge_base = KnowledgeBase.objects.filter(id=kb_id).first() if not knowledge_base: return Response({ 'code': 404, 'message': f'知识库不存在: {kb_id}', 'data': None }, status=status.HTTP_404_NOT_FOUND) knowledge_bases.append(knowledge_base) # 使用统一的权限检查方法 if not self.check_knowledge_base_permission(knowledge_base, user, 'read'): return Response({ 'code': 403, 'message': f'无权访问知识库: {knowledge_base.name}', 'data': None }, status=status.HTTP_403_FORBIDDEN) except Exception as e: return Response({ 'code': 400, 'message': f'处理知识库ID出错: {str(e)}', 'data': None }, status=status.HTTP_400_BAD_REQUEST) # 创建一个新的会话ID conversation_id = str(uuid.uuid4()) logger.info(f"创建新的会话ID: {conversation_id}") # 准备metadata (仍然保存知识库名称用于内部处理) metadata = { 'dataset_id_list': [str(id) for id in dataset_ids], 'dataset_names': [kb.name for kb in knowledge_bases] } return Response({ 'code': 200, 'message': '会话创建成功', 'data': { 'conversation_id': conversation_id, 'dataset_id_list': metadata['dataset_id_list'] } }) except Exception as e: logger.error(f"创建会话失败: {str(e)}") logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f'创建会话失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) def create(self, request): """创建聊天记录""" try: chat_service = ChatService() question_record, conversation_id, metadata, knowledge_bases, external_id_list = chat_service.create_chat_record( request.user, request.data, request.data.get('conversation_id') ) use_stream = request.data.get('stream', True) title = request.data.get('title', 'New chat') if use_stream: def stream_response(): answer_record = ChatHistory.objects.create( user=question_record.user, knowledge_base=knowledge_bases[0], conversation_id=conversation_id, title=title, parent_id=str(question_record.id), role='assistant', content="", metadata=metadata ) yield f"data: {json.dumps({'code': 200, 'message': '开始流式传输', 'data': {'id': str(answer_record.id), 'conversation_id': conversation_id, 'content': '', 'is_end': False}})}\n\n" full_content = "" for data in stream_chat_answer(conversation_id, request.data['question'], external_id_list, metadata): parsed_data = json.loads(data[5:-2]) # 移除"data: "和"\n\n" if parsed_data['code'] == 200 and 'content' in parsed_data['data']: content_part = parsed_data['data']['content'] full_content += content_part response_data = { 'code': 200, 'message': 'partial', 'data': { 'id': str(answer_record.id), 'conversation_id': conversation_id, 'title': title, 'content': content_part, 'is_end': parsed_data['data']['is_end'] } } yield f"data: {json.dumps(response_data)}\n\n" if parsed_data['data']['is_end']: answer_record.content = full_content.strip() answer_record.save() current_title = ChatHistory.objects.filter( conversation_id=conversation_id ).exclude( title__in=["New chat", "新对话", ""] ).values_list('title', flat=True).first() if current_title: title_updated = current_title else: try: generated_title = generate_conversation_title( request.data['question'], full_content.strip() ) if generated_title: ChatHistory.objects.filter( conversation_id=conversation_id ).update(title=generated_title) title_updated = generated_title else: title_updated = "新对话" except ExternalAPIError as e: logger.error(f"自动生成标题失败: {str(e)}") title_updated = "新对话" final_response = { 'code': 200, 'message': '完成', 'data': { 'id': str(answer_record.id), 'conversation_id': conversation_id, 'title': title_updated, 'dataset_id_list': metadata.get('dataset_id_list', []), 'dataset_names': metadata.get('dataset_names', []), 'role': 'assistant', 'content': full_content.strip(), 'created_at': answer_record.created_at.strftime('%Y-%m-%d %H:%M:%S'), 'is_end': True } } yield f"data: {json.dumps(final_response)}\n\n" break elif parsed_data['code'] != 200: yield data break if full_content: try: answer_record.content = full_content.strip() answer_record.save() except Exception as save_error: logger.error(f"保存部分内容失败: {str(save_error)}") response = StreamingHttpResponse( stream_response(), content_type='text/event-stream', status=status.HTTP_201_CREATED ) response['Cache-Control'] = 'no-cache, no-store' response['Connection'] = 'keep-alive' return response else: logger.info("使用非流式输出模式") try: answer = get_chat_answer(external_id_list, request.data['question']) except ExternalAPIError as e: logger.error(f"获取回答失败: {str(e)}") return Response({ 'code': 500, 'message': f'获取回答失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) if answer is None: return Response({ 'code': 500, 'message': '获取回答失败', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) answer_record = ChatHistory.objects.create( user=request.user, knowledge_base=knowledge_bases[0], conversation_id=conversation_id, title=title, parent_id=str(question_record.id), role='assistant', content=answer, metadata=metadata ) existing_records = ChatHistory.objects.filter(conversation_id=conversation_id) should_generate_title = not existing_records.exclude(id=question_record.id).exists() and (not title or title == 'New chat') if should_generate_title: try: generated_title = generate_conversation_title( request.data['question'], answer ) if generated_title: ChatHistory.objects.filter(conversation_id=conversation_id).update(title=generated_title) title = generated_title except ExternalAPIError as e: logger.error(f"自动生成标题失败: {str(e)}") return Response({ 'code': 200, 'message': '成功', 'data': { 'id': str(answer_record.id), 'conversation_id': conversation_id, 'title': title, 'dataset_id_list': metadata.get('dataset_id_list', []), 'dataset_names': metadata.get('dataset_names', []), 'role': 'assistant', 'content': answer, 'created_at': answer_record.created_at.strftime('%Y-%m-%d %H:%M:%S') } }, status=status.HTTP_201_CREATED) except ValueError as e: return Response({ 'code': 400, 'message': str(e), 'data': None }, status=status.HTTP_400_BAD_REQUEST) except Exception as e: logger.error(f"创建聊天记录失败: {str(e)}") import traceback logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f'创建聊天记录失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @action(detail=False, methods=['post']) def hit_test(self, request): """获取问题与知识库文档的匹配度""" try: data = request.data if 'question' not in data or 'dataset_id_list' not in data or not data['dataset_id_list']: return Response({ 'code': 400, 'message': '缺少必填字段: question 或 dataset_id_list', 'data': None }, status=status.HTTP_400_BAD_REQUEST) question = data['question'] dataset_ids = data['dataset_id_list'] if not isinstance(dataset_ids, list): try: dataset_ids = json.loads(dataset_ids) if not isinstance(dataset_ids, list): dataset_ids = [dataset_ids] except (json.JSONDecodeError, TypeError): dataset_ids = [dataset_ids] external_id_list = [] for kb_id in dataset_ids: kb = KnowledgeBase.objects.filter(id=kb_id).first() if not kb: return Response({ 'code': 404, 'message': f'知识库不存在: {kb_id}', 'data': None }, status=status.HTTP_404_NOT_FOUND) if not self.check_knowledge_base_permission(kb, request.user, 'read'): return Response({ 'code': 403, 'message': f'无权访问知识库: {kb.name}', 'data': None }, status=status.HTTP_403_FORBIDDEN) if kb.external_id: external_id_list.append(str(kb.external_id)) if not external_id_list: return Response({ 'code': 400, 'message': '没有有效的知识库external_id', 'data': None }, status=status.HTTP_400_BAD_REQUEST) all_documents = [] for dataset_id in external_id_list: try: doc_info = get_hit_test_documents(dataset_id, question) if doc_info: all_documents.extend(doc_info) except ExternalAPIError as e: logger.error(f"调用hit_test失败: 知识库ID={dataset_id}, 错误={str(e)}") continue # 宽松处理，跳过失败的知识库 all_documents = sorted(all_documents, key=lambda x: x.get('similarity', 0), reverse=True) return Response({ 'code': 200, 'message': '成功', 'data': { 'question': question, 'matched_documents': all_documents, 'total_count': len(all_documents) } }) except Exception as e: logger.error(f"hit_test接口调用失败: {str(e)}") import traceback logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f'hit_test接口调用失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) def _highlight_keyword(self, text, keyword): """高亮关键词""" if not keyword or not text: return text return text.replace(keyword, f'{keyword}') def update(self, request, pk=None): """更新聊天记录""" try: record = self.get_queryset().filter(id=pk).first() if not record: return Response({ 'code': 404, 'message': '记录不存在或无权限', 'data': None }, status=status.HTTP_404_NOT_FOUND) data = request.data updateable_fields = ['content', 'metadata'] if 'content' in data: record.content = data['content'] if 'metadata' in data: current_metadata = record.metadata or {} current_metadata.update(data['metadata']) record.metadata = current_metadata record.save() return Response({ 'code': 200, 'message': '更新成功', 'data': { 'id': str(record.id), 'conversation_id': record.conversation_id, 'role': record.role, 'content': record.content, 'metadata': record.metadata, 'updated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S') } }) except Exception as e: logger.error(f"更新聊天记录失败: {str(e)}") import traceback logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f'更新聊天记录失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) def destroy(self, request, pk=None): """删除聊天记录（软删除）""" try: record = self.get_queryset().filter(id=pk).first() if not record: return Response({ 'code': 404, 'message': '记录不存在或无权限', 'data': None }, status=status.HTTP_404_NOT_FOUND) record.soft_delete() return Response({ 'code': 200, 'message': '删除成功', 'data': None }) except Exception as e: logger.error(f"删除聊天记录失败: {str(e)}") import traceback logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f'删除聊天记录失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @action(detail=False, methods=['get']) def search(self, request): """搜索聊天记录""" try: keyword = request.query_params.get('keyword', '').strip() dataset_id = request.query_params.get('dataset_id') start_date = request.query_params.get('start_date') end_date = request.query_params.get('end_date') page = int(request.query_params.get('page', 1)) page_size = int(request.query_params.get('page_size', 10)) query = self.get_queryset() if keyword: query = query.filter( Q(content__icontains=keyword) | Q(knowledge_base__name__icontains=keyword) ) if dataset_id: knowledge_base = KnowledgeBase.objects.filter(id=dataset_id).first() if knowledge_base and not self.check_knowledge_base_permission(knowledge_base, request.user, 'read'): return Response({ 'code': 403, 'message': '无权访问该知识库', 'data': None }, status=status.HTTP_403_FORBIDDEN) query = query.filter(knowledge_base__id=dataset_id) if start_date: query = query.filter(created_at__gte=start_date) if end_date: query = query.filter(created_at__lte=end_date) total = query.count() start = (page - 1) * page_size end = start + page_size records = query.order_by('-created_at')[start:end] results = [ { 'id': str(record.id), 'conversation_id': record.conversation_id, 'dataset_id': str(record.knowledge_base.id), 'dataset_name': record.knowledge_base.name, 'role': record.role, 'content': record.content, 'created_at': record.created_at.strftime('%Y-%m-%d %H:%M:%S'), 'metadata': record.metadata, 'highlights': {'content': self._highlight_keyword(record.content, keyword)} if keyword else {} } for record in records ] return Response({ 'code': 200, 'message': '搜索成功', 'data': { 'total': total, 'page': page, 'page_size': page_size, 'results': results } }) except Exception as e: logger.error(f"搜索聊天记录失败: {str(e)}") import traceback logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f'搜索失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @action(detail=False, methods=['delete']) def delete_conversation(self, request): """通过conversation_id删除一组会话""" try: conversation_id = request.query_params.get('conversation_id') if not conversation_id: return Response({ 'code': 400, 'message': '缺少必要参数: conversation_id', 'data': None }, status=status.HTTP_400_BAD_REQUEST) records = self.get_queryset().filter(conversation_id=conversation_id) if not records.exists(): return Response({ 'code': 404, 'message': '未找到该会话或无权限访问', 'data': None }, status=status.HTTP_404_NOT_FOUND) records_count = records.count() for record in records: record.soft_delete() return Response({ 'code': 200, 'message': '删除成功', 'data': { 'conversation_id': conversation_id, 'deleted_count': records_count } }) except Exception as e: logger.error(f"删除会话失败: {str(e)}") import traceback logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f'删除会话失败: {str(e)}', 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @action(detail=False, methods=['get'], url_path='generate-conversation-title') def generate_conversation_title(self, request): """更新会话标题""" try: conversation_id = request.query_params.get('conversation_id') if not conversation_id: return Response({ 'code': 400, 'message': '缺少conversation_id参数', 'data': None }, status=status.HTTP_400_BAD_REQUEST) # 检查对话是否存在 messages = self.get_queryset().filter( conversation_id=conversation_id, is_deleted=False, user=request.user ).order_by('created_at') if not messages.exists(): return Response({ 'code': 404, 'message': '对话不存在或无权访问', 'data': None }, status=status.HTTP_404_NOT_FOUND) # 检查是否有自定义标题参数 custom_title = request.query_params.get('title') if not custom_title: return Response({ 'code': 400, 'message': '缺少title参数', 'data': None }, status=status.HTTP_400_BAD_REQUEST) # 更新所有相关记录的标题 ChatHistory.objects.filter( conversation_id=conversation_id, user=request.user ).update(title=custom_title) return Response({ 'code': 200, 'message': '更新会话标题成功', 'data': { 'conversation_id': conversation_id, 'title': custom_title } }) except Exception as e: logger.error(f"更新会话标题失败: {str(e)}") logger.error(traceback.format_exc()) return Response({ 'code': 500, 'message': f"更新会话标题失败: {str(e)}", 'data': None }, status=status.HTTP_500_INTERNAL_SERVER_ERROR)