310 lines
9.8 KiB
Python
Executable File
310 lines
9.8 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
|
||
from app.models.node_template import NodeTemplate
|
||
from datetime import datetime
|
||
import uuid
|
||
|
||
|
||
def generate_complex_agent_with_templates(db: Session, username: str = "admin"):
|
||
"""生成包含多个节点模板的复杂Agent"""
|
||
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()
|
||
|
||
# 查找可用的节点模板（公开的或用户自己的）
|
||
templates = db.query(NodeTemplate).filter(
|
||
(NodeTemplate.is_public == True) | (NodeTemplate.user_id == user.id)
|
||
).limit(10).all()
|
||
|
||
if len(templates) < 3:
|
||
print(f"⚠️ 警告: 只找到 {len(templates)} 个节点模板，建议至少3个")
|
||
if len(templates) == 0:
|
||
print("❌ 没有可用的节点模板，请先创建节点模板")
|
||
return
|
||
|
||
print(f"📋 找到 {len(templates)} 个可用节点模板:")
|
||
for i, template in enumerate(templates[:5], 1):
|
||
print(f" {i}. {template.name} (ID: {template.id})")
|
||
print()
|
||
|
||
# 选择3-5个模板用于创建复杂Agent
|
||
selected_templates = templates[:min(5, len(templates))]
|
||
template_ids = [t.id for t in selected_templates]
|
||
|
||
# 生成复杂的工作流配置
|
||
# 工作流结构：开始 -> 模板1 -> 条件判断 -> 模板2/模板3 -> 转换 -> 模板4 -> 结束
|
||
nodes = []
|
||
edges = []
|
||
|
||
# 开始节点
|
||
start_node = {
|
||
"id": "start-1",
|
||
"type": "start",
|
||
"position": {"x": 100, "y": 200},
|
||
"data": {"label": "开始"}
|
||
}
|
||
nodes.append(start_node)
|
||
|
||
# 第一个模板节点
|
||
if len(template_ids) > 0:
|
||
template1_node = {
|
||
"id": "template-1",
|
||
"type": "template",
|
||
"position": {"x": 300, "y": 200},
|
||
"data": {
|
||
"label": selected_templates[0].name,
|
||
"template_id": template_ids[0],
|
||
"provider": selected_templates[0].provider,
|
||
"model": selected_templates[0].model,
|
||
"temperature": selected_templates[0].temperature,
|
||
"max_tokens": selected_templates[0].max_tokens,
|
||
"prompt": selected_templates[0].prompt
|
||
}
|
||
}
|
||
nodes.append(template1_node)
|
||
edges.append({
|
||
"id": "e1",
|
||
"source": "start-1",
|
||
"target": "template-1"
|
||
})
|
||
|
||
# 条件判断节点
|
||
condition_node = {
|
||
"id": "condition-1",
|
||
"type": "condition",
|
||
"position": {"x": 500, "y": 200},
|
||
"data": {
|
||
"label": "判断处理结果",
|
||
"condition": "{{result}} contains '需要' or {{result}} contains '重要'"
|
||
}
|
||
}
|
||
nodes.append(condition_node)
|
||
if len(template_ids) > 0:
|
||
edges.append({
|
||
"id": "e2",
|
||
"source": "template-1",
|
||
"target": "condition-1"
|
||
})
|
||
|
||
# 条件分支：如果满足条件，使用模板2；否则使用模板3
|
||
if len(template_ids) > 1:
|
||
template2_node = {
|
||
"id": "template-2",
|
||
"type": "template",
|
||
"position": {"x": 700, "y": 100},
|
||
"data": {
|
||
"label": selected_templates[1].name,
|
||
"template_id": template_ids[1],
|
||
"provider": selected_templates[1].provider,
|
||
"model": selected_templates[1].model,
|
||
"temperature": selected_templates[1].temperature,
|
||
"max_tokens": selected_templates[1].max_tokens,
|
||
"prompt": selected_templates[1].prompt
|
||
}
|
||
}
|
||
nodes.append(template2_node)
|
||
edges.append({
|
||
"id": "e3",
|
||
"source": "condition-1",
|
||
"target": "template-2",
|
||
"sourceHandle": "true"
|
||
})
|
||
|
||
if len(template_ids) > 2:
|
||
template3_node = {
|
||
"id": "template-3",
|
||
"type": "template",
|
||
"position": {"x": 700, "y": 300},
|
||
"data": {
|
||
"label": selected_templates[2].name,
|
||
"template_id": template_ids[2],
|
||
"provider": selected_templates[2].provider,
|
||
"model": selected_templates[2].model,
|
||
"temperature": selected_templates[2].temperature,
|
||
"max_tokens": selected_templates[2].max_tokens,
|
||
"prompt": selected_templates[2].prompt
|
||
}
|
||
}
|
||
nodes.append(template3_node)
|
||
edges.append({
|
||
"id": "e4",
|
||
"source": "condition-1",
|
||
"target": "template-3",
|
||
"sourceHandle": "false"
|
||
})
|
||
|
||
# 数据转换节点
|
||
transform_node = {
|
||
"id": "transform-1",
|
||
"type": "transform",
|
||
"position": {"x": 900, "y": 200},
|
||
"data": {
|
||
"label": "数据转换",
|
||
"mode": "mapping",
|
||
"mapping": {
|
||
"final_result": "{{result}}",
|
||
"processed": "true"
|
||
}
|
||
}
|
||
}
|
||
nodes.append(transform_node)
|
||
|
||
# 连接条件分支到转换节点
|
||
if len(template_ids) > 1:
|
||
edges.append({
|
||
"id": "e5",
|
||
"source": "template-2",
|
||
"target": "transform-1"
|
||
})
|
||
if len(template_ids) > 2:
|
||
edges.append({
|
||
"id": "e6",
|
||
"source": "template-3",
|
||
"target": "transform-1"
|
||
})
|
||
|
||
# 最后一个模板节点（如果还有模板）
|
||
if len(template_ids) > 3:
|
||
template4_node = {
|
||
"id": "template-4",
|
||
"type": "template",
|
||
"position": {"x": 1100, "y": 200},
|
||
"data": {
|
||
"label": selected_templates[3].name,
|
||
"template_id": template_ids[3],
|
||
"provider": selected_templates[3].provider,
|
||
"model": selected_templates[3].model,
|
||
"temperature": selected_templates[3].temperature,
|
||
"max_tokens": selected_templates[3].max_tokens,
|
||
"prompt": selected_templates[3].prompt
|
||
}
|
||
}
|
||
nodes.append(template4_node)
|
||
edges.append({
|
||
"id": "e7",
|
||
"source": "transform-1",
|
||
"target": "template-4"
|
||
})
|
||
# 结束节点连接到最后一个模板
|
||
end_node = {
|
||
"id": "end-1",
|
||
"type": "end",
|
||
"position": {"x": 1300, "y": 200},
|
||
"data": {"label": "结束"}
|
||
}
|
||
nodes.append(end_node)
|
||
edges.append({
|
||
"id": "e8",
|
||
"source": "template-4",
|
||
"target": "end-1"
|
||
})
|
||
else:
|
||
# 如果没有更多模板，转换节点直接连接到结束节点
|
||
end_node = {
|
||
"id": "end-1",
|
||
"type": "end",
|
||
"position": {"x": 1100, "y": 200},
|
||
"data": {"label": "结束"}
|
||
}
|
||
nodes.append(end_node)
|
||
edges.append({
|
||
"id": "e7",
|
||
"source": "transform-1",
|
||
"target": "end-1"
|
||
})
|
||
|
||
# 检查是否已存在同名Agent
|
||
agent_name = "复杂模板Agent（多节点）"
|
||
existing = db.query(Agent).filter(
|
||
Agent.name == agent_name,
|
||
Agent.user_id == user.id
|
||
).first()
|
||
|
||
if existing:
|
||
print(f"⚠️ Agent '{agent_name}' 已存在，将更新它...")
|
||
existing.workflow_config = {"nodes": nodes, "edges": edges}
|
||
existing.description = f"包含 {len(selected_templates)} 个节点模板的复杂Agent，支持条件分支和数据转换"
|
||
existing.updated_at = datetime.now()
|
||
agent = existing
|
||
else:
|
||
# 创建Agent
|
||
agent = Agent(
|
||
name=agent_name,
|
||
description=f"包含 {len(selected_templates)} 个节点模板的复杂Agent，支持条件分支和数据转换",
|
||
workflow_config={"nodes": nodes, "edges": edges},
|
||
status="published",
|
||
user_id=user.id,
|
||
version=1
|
||
)
|
||
db.add(agent)
|
||
|
||
try:
|
||
db.commit()
|
||
db.refresh(agent)
|
||
print()
|
||
print("=" * 60)
|
||
print(f"✅ 成功创建/更新复杂Agent: {agent.name}")
|
||
print(f" ID: {agent.id}")
|
||
print(f" 状态: {agent.status}")
|
||
print(f" 节点数量: {len(nodes)}")
|
||
print(f" 连接数量: {len(edges)}")
|
||
print(f" 使用的模板节点: {len(selected_templates)}")
|
||
print("=" * 60)
|
||
print()
|
||
print("工作流结构:")
|
||
print(" 开始 -> 模板1 -> 条件判断 -> 模板2/模板3 -> 数据转换 -> 模板4 -> 结束")
|
||
print()
|
||
return agent
|
||
except Exception as e:
|
||
db.rollback()
|
||
print()
|
||
print("=" * 60)
|
||
print(f"❌ 创建失败: {e}")
|
||
print("=" * 60)
|
||
import traceback
|
||
traceback.print_exc()
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="生成包含多个节点模板的复杂Agent")
|
||
parser.add_argument(
|
||
"--username",
|
||
type=str,
|
||
default="admin",
|
||
help="创建Agent的用户名（默认: admin）"
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
db = SessionLocal()
|
||
try:
|
||
generate_complex_agent_with_templates(db, username=args.username)
|
||
except Exception as e:
|
||
print(f"❌ 发生错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
finally:
|
||
db.close()
|