PlantAI-dev/example_task_with_wechat.py

#!/usr/bin/env python3
"""
示例：在 bbav20311 任务中集成微信通知

这个示例展示了如何在基因家族分析任务中
添加微信通知功能
"""

import asyncio
import time
from datetime import datetime
from pathlib import Path

# 导入微信通知模块
from notify_integration import notify_task_start, notify_task_complete


# 用户的微信ID（需要替换为真实ID）
USER_WECHAT_ID = "wxid_xxxxxxxxxxxxxxxx"  # 替换为你的微信ID


async def run_identify_with_notification(prefix: str, user_wechat_id: str):
    """
    运行基因家族鉴定任务，并发送微信通知
    
    使用示例:
        await run_identify_with_notification("Ft.SPL", "wxid_xxx")
    """
    task_id = f"identify_{prefix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    start_time = time.time()
    
    print(f"🚀 开始任务: {task_id}")
    
    # 1. 通知任务开始
    notify_task_start(user_wechat_id, task_id, f"基因家族鉴定 ({prefix})")
    
    try:
        # 2. 执行实际任务（这里用模拟代替）
        print(f"   正在分析 {prefix} 基因家族...")
        
        # 模拟任务执行
        await asyncio.sleep(5)
        
        # 3. 任务完成，通知用户
        duration = f"{int(time.time() - start_time)}秒"
        output_files = [
            f"{prefix}.identify.gene.list.txt",
            f"{prefix}.identify.protein.fa",
            f"{prefix}.identify.report.pdf"
        ]
        
        notify_task_complete(
            user_wechat_id=user_wechat_id,
            task_id=task_id,
            task_type=f"基因家族鉴定 ({prefix})",
            output_files=output_files,
            duration=duration
        )
        
        print(f"✅ 任务完成: {task_id}")
        
    except Exception as e:
        # 4. 任务失败，通知用户
        duration = f"{int(time.time() - start_time)}秒"
        notify_task_complete(
            user_wechat_id=user_wechat_id,
            task_id=task_id,
            task_type=f"基因家族鉴定 ({prefix})",
            output_files=[],
            duration=duration,
            error=str(e)
        )
        
        print(f"❌ 任务失败: {task_id}, 错误: {e}")


async def run_pipeline_with_notifications(
    prefix: str, 
    stages: list,
    user_wechat_id: str
):
    """
    运行多阶段 pipeline，每个阶段都发送通知
    
    使用示例:
        await run_pipeline_with_notifications(
            "Ft.SPL",
            ["identify", "phylogeny", "motif"],
            "wxid_xxx"
        )
    """
    pipeline_id = f"pipeline_{prefix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    start_time = time.time()
    
    print(f"🚀 开始 Pipeline: {pipeline_id}")
    print(f"   阶段: {', '.join(stages)}")
    
    # 通知 pipeline 开始
    notify_task_start(
        user_wechat_id, 
        pipeline_id, 
        f"综合分析 Pipeline ({prefix})"
    )
    
    completed_stages = []
    failed_stages = []
    
    for stage in stages:
        stage_start = time.time()
        
        try:
            print(f"\n   执行阶段: {stage}")
            
            # 模拟阶段执行
            await asyncio.sleep(3)
            
            completed_stages.append(stage)
            
        except Exception as e:
            failed_stages.append((stage, str(e)))
            print(f"   阶段失败: {stage}, 错误: {e}")
    
    # Pipeline 完成通知
    duration = f"{int(time.time() - start_time)}秒"
    
    if failed_stages:
        error_msg = f"完成 {len(completed_stages)}/{len(stages)} 个阶段\n"
        error_msg += f"失败阶段: {', '.join([s for s, _ in failed_stages])}"
        
        notify_task_complete(
            user_wechat_id=user_wechat_id,
            task_id=pipeline_id,
            task_type=f"综合分析 Pipeline ({prefix})",
            output_files=[f"{prefix}.{s}.result.txt" for s in completed_stages],
            duration=duration,
            error=error_msg
        )
    else:
        notify_task_complete(
            user_wechat_id=user_wechat_id,
            task_id=pipeline_id,
            task_type=f"综合分析 Pipeline ({prefix})",
            output_files=[f"{prefix}.{s}.result.txt" for s in completed_stages],
            duration=duration
        )
    
    print(f"\n✅ Pipeline 完成: {pipeline_id}")


# ==================== 测试 ====================

async def main():
    """测试示例"""
    print("="*60)
    print("🧪 微信通知集成示例")
    print("="*60)
    
    if USER_WECHAT_ID == "wxid_xxxxxxxxxxxxxxxx":
        print("\n⚠️  请修改 USER_WECHAT_ID 为你的真实微信ID")
        print("   文件位置: example_task_with_wechat.py")
        return
    
    # 测试单任务
    print("\n📋 测试1: 单任务通知")
    await run_identify_with_notification("Ft.SPL", USER_WECHAT_ID)
    
    # 测试 Pipeline
    print("\n📋 测试2: Pipeline 通知")
    await run_pipeline_with_notifications(
        "Ft.WRKY",
        ["identify", "phylogeny", "motif"],
        USER_WECHAT_ID
    )
    
    print("\n" + "="*60)
    print("✅ 所有测试完成")
    print("="*60)


if __name__ == "__main__":
    asyncio.run(main())