#!/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()