310 lines
9.8 KiB
Python
310 lines
9.8 KiB
Python
|
#!/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()
|