aiagent/backend/scripts/generate_test_agent.py

#!/usr/bin/env python3
"""
生成测试Agent - 演示左右连接
这是一个测试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 datetime import datetime
import uuid


def generate_test_agent(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()
    
    # 生成测试工作流配置
    # 工作流结构（横向排列，使用左右连接）：
    # 开始 → 处理1 → 处理2 → 处理3 → 结束
    
    nodes = []
    edges = []
    
    # 节点横向排列，Y坐标相同，X坐标递增
    base_y = 300
    x_spacing = 250
    
    # 1. 开始节点
    start_node = {
        "id": "start-1",
        "type": "start",
        "position": {"x": 50, "y": base_y},
        "data": {
            "label": "开始",
            "output_format": "json"
        }
    }
    nodes.append(start_node)
    
    # 2. 处理节点1（LLM）
    process1_node = {
        "id": "llm-process1",
        "type": "llm",
        "position": {"x": 50 + x_spacing, "y": base_y},
        "data": {
            "label": "处理1",
            "provider": "deepseek",
            "model": "deepseek-chat",
            "temperature": "0.3",
            "max_tokens": "1000",
            "prompt": """你是一个测试助手。请简单处理用户输入。

用户输入：{{query}}

请输出：已处理用户输入："{{query}}"
"""
        }
    }
    nodes.append(process1_node)
    
    # 3. 处理节点2（Transform）
    process2_node = {
        "id": "transform-process2",
        "type": "transform",
        "position": {"x": 50 + x_spacing * 2, "y": base_y},
        "data": {
            "label": "处理2",
            "mode": "merge",
            "mapping": {
                "original_input": "{{query}}",
                "processed_result": "{{output}}"
            }
        }
    }
    nodes.append(process2_node)
    
    # 4. 处理节点3（LLM）
    process3_node = {
        "id": "llm-process3",
        "type": "llm",
        "position": {"x": 50 + x_spacing * 3, "y": base_y},
        "data": {
            "label": "处理3",
            "provider": "deepseek",
            "model": "deepseek-chat",
            "temperature": "0.3",
            "max_tokens": "1000",
            "prompt": """请对处理结果进行总结。

原始输入：{{original_input}}
处理结果：{{processed_result}}

请输出一个简洁的总结。"""
        }
    }
    nodes.append(process3_node)
    
    # 5. 结束节点
    end_node = {
        "id": "end-1",
        "type": "end",
        "position": {"x": 50 + x_spacing * 4, "y": base_y},
        "data": {
            "label": "结束",
            "description": "返回最终结果"
        }
    }
    nodes.append(end_node)
    
    # 创建连接（使用左右连接）
    # 注意：使用 sourceHandle='right' 和 targetHandle='left' 来指定左右连接
    edges.append({
        "id": "e1",
        "source": "start-1",
        "target": "llm-process1",
        "sourceHandle": "right",  # 从开始节点的右侧连接
        "targetHandle": "left"    # 连接到处理1节点的左侧
    })
    
    edges.append({
        "id": "e2",
        "source": "llm-process1",
        "target": "transform-process2",
        "sourceHandle": "right",  # 从处理1节点的右侧连接
        "targetHandle": "left"    # 连接到处理2节点的左侧
    })
    
    edges.append({
        "id": "e3",
        "source": "transform-process2",
        "target": "llm-process3",
        "sourceHandle": "right",  # 从处理2节点的右侧连接
        "targetHandle": "left"    # 连接到处理3节点的左侧
    })
    
    edges.append({
        "id": "e4",
        "source": "llm-process3",
        "target": "end-1",
        "sourceHandle": "right",  # 从处理3节点的右侧连接
        "targetHandle": "left"    # 连接到结束节点的左侧
    })
    
    # 创建或更新Agent
    workflow_config = {
        "nodes": nodes,
        "edges": edges
    }
    
    agent = db.query(Agent).filter(
        Agent.name == "测试Agent（左右连接）",
        Agent.user_id == user.id
    ).first()
    
    if agent:
        agent.workflow_config = workflow_config
        agent.description = "测试Agent，演示节点左右连接功能。工作流横向排列，使用左右连接点。"
        agent.updated_at = datetime.now()
        agent.status = "published"
        print("⚠️  Agent '测试Agent（左右连接）' 已存在，将更新它...")
    else:
        agent = Agent(
            id=str(uuid.uuid4()),
            name="测试Agent（左右连接）",
            description="测试Agent，演示节点左右连接功能。工作流横向排列，使用左右连接点。",
            workflow_config=workflow_config,
            status="published",
            user_id=user.id,
            version=1
        )
        db.add(agent)
    
    try:
        db.commit()
        db.refresh(agent)
        print()
        print("✅ Agent创建/更新成功！")
        print()
        print(f"📋 Agent信息：")
        print(f"   - ID: {agent.id}")
        print(f"   - 名称: {agent.name}")
        print(f"   - 状态: {agent.status} (已发布，可直接使用)")
        print(f"   - 版本: {agent.version}")
        print(f"   - 节点数: {len(nodes)}")
        print(f"   - 连接数: {len(edges)}")
        print()
        print("🎯 工作流特点：")
        print("   ✅ 节点横向排列")
        print("   ✅ 使用左右连接点（sourceHandle='right', targetHandle='left'）")
        print("   ✅ 工作流：开始 → 处理1 → 处理2 → 处理3 → 结束")
        print()
        print("💡 使用提示：")
        print("   1. 在Agent管理页面找到 '测试Agent（左右连接）'")
        print("   2. 点击 '使用' 按钮开始使用")
        print("   3. 输入任意文本进行测试")
        print("   4. 在工作流设计器中查看左右连接的视觉效果")
        print()
        return agent
    except Exception as e:
        db.rollback()
        print(f"❌ 创建Agent失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return None


def main():
    """主函数"""
    db = SessionLocal()
    try:
        generate_test_agent(db, username="admin")
    finally:
        db.close()


if __name__ == "__main__":
    main()