55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
from django.core.management.base import BaseCommand
|
|
from django.contrib.auth import get_user_model
|
|
from user_management.models import ChatHistory
|
|
from django.utils import timezone
|
|
import uuid
|
|
|
|
User = get_user_model()
|
|
|
|
class Command(BaseCommand):
|
|
help = '创建聊天测试数据'
|
|
|
|
def handle(self, *args, **kwargs):
|
|
try:
|
|
# 获取测试用户
|
|
user = User.objects.first()
|
|
if not user:
|
|
self.stdout.write(self.style.ERROR('没有找到用户，请先创建用户'))
|
|
return
|
|
|
|
# 创建3个不同的对话
|
|
for i in range(3):
|
|
conversation_id = f"conv_{uuid.uuid4().hex[:10]}"
|
|
dataset_id = f"dataset_{i+1}"
|
|
dataset_name = f"测试数据集_{i+1}"
|
|
|
|
# 每个对话创建3轮问答
|
|
for j in range(3):
|
|
# 用户问题
|
|
ChatHistory.objects.create(
|
|
user=user,
|
|
dataset_id=dataset_id,
|
|
dataset_name=dataset_name,
|
|
conversation_id=conversation_id,
|
|
role='user',
|
|
content=f"这是第{i+1}个数据集的第{j+1}个问题？",
|
|
question=f"这是第{i+1}个数据集的第{j+1}个问题？",
|
|
answer="",
|
|
)
|
|
|
|
# AI回答
|
|
ChatHistory.objects.create(
|
|
user=user,
|
|
dataset_id=dataset_id,
|
|
dataset_name=dataset_name,
|
|
conversation_id=conversation_id,
|
|
role='assistant',
|
|
content=f"这是第{i+1}个数据集的第{j+1}个回答。",
|
|
question="",
|
|
answer=f"这是第{i+1}个数据集的第{j+1}个回答。",
|
|
)
|
|
|
|
self.stdout.write(self.style.SUCCESS('成功创建测试聊天数据'))
|
|
|
|
except Exception as e:
|
|
self.stdout.write(self.style.ERROR(f'创建测试数据失败: {str(e)}')) |