340 lines
11 KiB
Python
Executable File
340 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
生成Agent假数据
|
||
"""
|
||
import sys
|
||
import os
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
from sqlalchemy.orm import Session
|
||
from app.core.database import SessionLocal
|
||
from app.models.agent import Agent
|
||
from app.models.user import User
|
||
import random
|
||
from datetime import datetime, timedelta
|
||
|
||
# 假数据模板
|
||
AGENT_TEMPLATES = [
|
||
{
|
||
"name": "智能客服助手",
|
||
"description": "自动处理客户咨询，提供7x24小时在线服务，支持多轮对话和智能转接",
|
||
"status": "running"
|
||
},
|
||
{
|
||
"name": "数据分析Agent",
|
||
"description": "自动分析业务数据，生成可视化报表，支持多种数据源接入",
|
||
"status": "published"
|
||
},
|
||
{
|
||
"name": "内容生成助手",
|
||
"description": "基于LLM的内容创作工具，支持文章、报告、营销文案等多种类型",
|
||
"status": "published"
|
||
},
|
||
{
|
||
"name": "邮件自动回复",
|
||
"description": "智能识别邮件内容，自动生成回复建议，提高工作效率",
|
||
"status": "draft"
|
||
},
|
||
{
|
||
"name": "代码审查Agent",
|
||
"description": "自动审查代码质量，检测潜在bug和安全漏洞，提供改进建议",
|
||
"status": "published"
|
||
},
|
||
{
|
||
"name": "会议纪要生成器",
|
||
"description": "自动记录会议内容，提取关键信息，生成结构化会议纪要",
|
||
"status": "running"
|
||
},
|
||
{
|
||
"name": "翻译助手",
|
||
"description": "支持多语言翻译，保持上下文连贯性，适用于文档和对话翻译",
|
||
"status": "published"
|
||
},
|
||
{
|
||
"name": "知识库问答",
|
||
"description": "基于企业知识库的智能问答系统，快速检索和回答专业问题",
|
||
"status": "running"
|
||
},
|
||
{
|
||
"name": "文档摘要生成",
|
||
"description": "自动提取文档关键信息，生成简洁准确的摘要，支持多种文档格式",
|
||
"status": "draft"
|
||
},
|
||
{
|
||
"name": "情感分析Agent",
|
||
"description": "分析文本情感倾向，监控用户反馈，识别潜在问题",
|
||
"status": "published"
|
||
},
|
||
{
|
||
"name": "任务调度助手",
|
||
"description": "智能分配和调度任务，优化资源利用，提高团队协作效率",
|
||
"status": "stopped"
|
||
},
|
||
{
|
||
"name": "API集成Agent",
|
||
"description": "连接多个外部API，实现数据同步和业务流程自动化",
|
||
"status": "published"
|
||
},
|
||
{
|
||
"name": "报表自动化",
|
||
"description": "定时生成各类业务报表，自动发送给相关人员，支持多种数据源",
|
||
"status": "running"
|
||
},
|
||
{
|
||
"name": "智能推荐系统",
|
||
"description": "基于用户行为和偏好，提供个性化推荐，提升用户体验",
|
||
"status": "draft"
|
||
},
|
||
{
|
||
"name": "异常检测Agent",
|
||
"description": "实时监控系统运行状态，自动检测异常并发送告警通知",
|
||
"status": "published"
|
||
}
|
||
]
|
||
|
||
|
||
def generate_workflow_config(agent_type: str) -> dict:
|
||
"""生成工作流配置"""
|
||
# 根据不同的Agent类型生成不同的工作流配置
|
||
base_configs = {
|
||
"客服": {
|
||
"nodes": [
|
||
{
|
||
"id": "start-1",
|
||
"type": "start",
|
||
"position": {"x": 100, "y": 100},
|
||
"data": {"label": "开始"}
|
||
},
|
||
{
|
||
"id": "llm-1",
|
||
"type": "llm",
|
||
"position": {"x": 300, "y": 100},
|
||
"data": {
|
||
"label": "LLM处理",
|
||
"model": "gpt-4",
|
||
"prompt": "你是一个专业的客服助手，请友好地回答用户的问题。"
|
||
}
|
||
},
|
||
{
|
||
"id": "condition-1",
|
||
"type": "condition",
|
||
"position": {"x": 500, "y": 100},
|
||
"data": {
|
||
"label": "判断是否需要转人工",
|
||
"condition": "{{需要人工}}"
|
||
}
|
||
},
|
||
{
|
||
"id": "end-1",
|
||
"type": "end",
|
||
"position": {"x": 700, "y": 100},
|
||
"data": {"label": "结束"}
|
||
}
|
||
],
|
||
"edges": [
|
||
{"id": "e1", "source": "start-1", "target": "llm-1"},
|
||
{"id": "e2", "source": "llm-1", "target": "condition-1"},
|
||
{"id": "e3", "source": "condition-1", "target": "end-1"}
|
||
]
|
||
},
|
||
"数据分析": {
|
||
"nodes": [
|
||
{
|
||
"id": "start-1",
|
||
"type": "start",
|
||
"position": {"x": 100, "y": 100},
|
||
"data": {"label": "开始"}
|
||
},
|
||
{
|
||
"id": "data-source-1",
|
||
"type": "data_source",
|
||
"position": {"x": 300, "y": 100},
|
||
"data": {
|
||
"label": "数据源",
|
||
"source_type": "database"
|
||
}
|
||
},
|
||
{
|
||
"id": "process-1",
|
||
"type": "process",
|
||
"position": {"x": 500, "y": 100},
|
||
"data": {
|
||
"label": "数据处理",
|
||
"operation": "aggregate"
|
||
}
|
||
},
|
||
{
|
||
"id": "llm-1",
|
||
"type": "llm",
|
||
"position": {"x": 700, "y": 100},
|
||
"data": {
|
||
"label": "生成分析报告",
|
||
"model": "gpt-4"
|
||
}
|
||
},
|
||
{
|
||
"id": "end-1",
|
||
"type": "end",
|
||
"position": {"x": 900, "y": 100},
|
||
"data": {"label": "结束"}
|
||
}
|
||
],
|
||
"edges": [
|
||
{"id": "e1", "source": "start-1", "target": "data-source-1"},
|
||
{"id": "e2", "source": "data-source-1", "target": "process-1"},
|
||
{"id": "e3", "source": "process-1", "target": "llm-1"},
|
||
{"id": "e4", "source": "llm-1", "target": "end-1"}
|
||
]
|
||
},
|
||
"默认": {
|
||
"nodes": [
|
||
{
|
||
"id": "start-1",
|
||
"type": "start",
|
||
"position": {"x": 100, "y": 100},
|
||
"data": {"label": "开始"}
|
||
},
|
||
{
|
||
"id": "llm-1",
|
||
"type": "llm",
|
||
"position": {"x": 300, "y": 100},
|
||
"data": {
|
||
"label": "LLM处理",
|
||
"model": "gpt-3.5-turbo",
|
||
"prompt": "请处理用户请求。"
|
||
}
|
||
},
|
||
{
|
||
"id": "end-1",
|
||
"type": "end",
|
||
"position": {"x": 500, "y": 100},
|
||
"data": {"label": "结束"}
|
||
}
|
||
],
|
||
"edges": [
|
||
{"id": "e1", "source": "start-1", "target": "llm-1"},
|
||
{"id": "e2", "source": "llm-1", "target": "end-1"}
|
||
]
|
||
}
|
||
}
|
||
|
||
# 根据名称判断类型
|
||
if "客服" in agent_type:
|
||
return base_configs["客服"]
|
||
elif "数据" in agent_type:
|
||
return base_configs["数据分析"]
|
||
else:
|
||
return base_configs["默认"]
|
||
|
||
|
||
def generate_fake_agents(db: Session, username: str = "admin", count: int = None):
|
||
"""生成假数据"""
|
||
print("=" * 60)
|
||
print("生成Agent假数据")
|
||
print("=" * 60)
|
||
print()
|
||
|
||
# 查找用户
|
||
user = db.query(User).filter(User.username == username).first()
|
||
if not user:
|
||
print(f"❌ 未找到用户 '{username}'，请先创建该用户")
|
||
return
|
||
|
||
print(f"✅ 找到用户: {user.username} (ID: {user.id})")
|
||
print()
|
||
|
||
# 确定要生成的数量
|
||
if count is None:
|
||
count = len(AGENT_TEMPLATES)
|
||
else:
|
||
count = min(count, len(AGENT_TEMPLATES))
|
||
|
||
print(f"📝 将生成 {count} 个Agent...")
|
||
print()
|
||
|
||
created_count = 0
|
||
skipped_count = 0
|
||
|
||
for i, template in enumerate(AGENT_TEMPLATES[:count]):
|
||
# 检查是否已存在同名Agent
|
||
existing = db.query(Agent).filter(
|
||
Agent.name == template["name"],
|
||
Agent.user_id == user.id
|
||
).first()
|
||
|
||
if existing:
|
||
print(f"⏭️ 跳过: {template['name']} (已存在)")
|
||
skipped_count += 1
|
||
continue
|
||
|
||
# 生成工作流配置
|
||
workflow_config = generate_workflow_config(template["name"])
|
||
|
||
# 随机生成创建时间（过去30天内）
|
||
days_ago = random.randint(0, 30)
|
||
created_at = datetime.now() - timedelta(days=days_ago)
|
||
updated_at = created_at + timedelta(hours=random.randint(1, 72))
|
||
|
||
# 创建Agent
|
||
agent = Agent(
|
||
name=template["name"],
|
||
description=template["description"],
|
||
workflow_config=workflow_config,
|
||
status=template["status"],
|
||
user_id=user.id,
|
||
version=random.randint(1, 5),
|
||
created_at=created_at,
|
||
updated_at=updated_at
|
||
)
|
||
|
||
db.add(agent)
|
||
created_count += 1
|
||
print(f"✅ 创建: {template['name']} (状态: {template['status']})")
|
||
|
||
# 提交事务
|
||
try:
|
||
db.commit()
|
||
print()
|
||
print("=" * 60)
|
||
print(f"✅ 成功生成 {created_count} 个Agent")
|
||
if skipped_count > 0:
|
||
print(f"⏭️ 跳过 {skipped_count} 个已存在的Agent")
|
||
print("=" * 60)
|
||
except Exception as e:
|
||
db.rollback()
|
||
print()
|
||
print("=" * 60)
|
||
print(f"❌ 生成失败: {e}")
|
||
print("=" * 60)
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="生成Agent假数据")
|
||
parser.add_argument(
|
||
"--username",
|
||
type=str,
|
||
default="admin",
|
||
help="创建Agent的用户名（默认: admin）"
|
||
)
|
||
parser.add_argument(
|
||
"--count",
|
||
type=int,
|
||
default=None,
|
||
help="要生成的Agent数量（默认: 生成所有模板）"
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
db = SessionLocal()
|
||
try:
|
||
generate_fake_agents(db, username=args.username, count=args.count)
|
||
except Exception as e:
|
||
print(f"❌ 发生错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
finally:
|
||
db.close()
|