lambdagent/agentexample/agent67/tools/research_workflow.py

"""
agent67 ResearchWorkflow 工具 — 一键科研流程

将 research skill pack 的 4 阶段 pipeline 封装为 agent67 的一个工具。
agent67 调用一次 ResearchWorkflow，内部自动串联:
  读论文 → 复现 → 实验 → 记录

Lambda 语义:
  ResearchWorkflow = λpaper. notebook(experiment(reproduce(read(paper))))

用法 (agent67 工具调用):
  {"action": "ResearchWorkflow", "input": {"paper": "Attention Is All You Need"}}
  {"action": "ResearchWorkflow", "input": {"paper": "https://arxiv.org/abs/1706.03762"}}
  {"action": "ResearchWorkflow", "input": {"paper": "transformer", "workspace": "./research/transformer"}}
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from lambdagent.core import Context


def research_workflow(input_str: str) -> str:
    """
    执行完整科研流程: 读论文→复现→实验→记录

    输入格式:
      JSON: {"paper": "论文标题/URL", "workspace": "保存路径(可选)", "skip": ["reproduce"](可选)}
      纯文本: 论文标题或URL
    """
    # 解析输入
    if isinstance(input_str, str):
        try:
            data = json.loads(input_str)
            paper = data.get("paper", data.get("query", input_str))
            workspace = data.get("workspace", "")
            skip_phases = data.get("skip", [])
        except (json.JSONDecodeError, AttributeError):
            paper = input_str.strip()
            workspace = ""
            skip_phases = []
    else:
        paper = str(input_str)
        workspace = ""
        skip_phases = []

    if not paper:
        return "[ResearchWorkflow] 请提供论文标题或URL"

    # 设置工作空间
    if not workspace:
        safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in paper[:50])
        workspace = f"./research/{safe_name}"

    os.makedirs(workspace, exist_ok=True)

    results = {}
    t0 = time.time()

    # ── Phase 1: 读论文 ──
    if "read" not in skip_phases:
        print(f"  📖 [1/4] 读论文: {paper[:60]}...")
        try:
            reader = _get_skill("paper-reader")
            if reader:
                paper_info = reader.apply(paper, Context())
                results["paper_info"] = str(paper_info)
                # 保存到工作空间
                _save_phase(workspace, "01_paper_reading", paper_info)
                print(f"  ✅ [1/4] 论文分析完成")
            else:
                # fallback: 简化版读论文 (用内置工具)
                results["paper_info"] = _fallback_read_paper(paper)
                print(f"  ✅ [1/4] 论文分析完成 (fallback)")
        except Exception as e:
            results["paper_info"] = f"[读论文失败: {e}]"
            print(f"  ⚠️  [1/4] 读论文失败: {e}")
    else:
        results["paper_info"] = "[跳过]"
        print(f"  ⏭️  [1/4] 跳过读论文")

    # ── Phase 2: 复现 ──
    if "reproduce" not in skip_phases:
        print(f"  🔬 [2/4] 代码复现...")
        try:
            reproducer = _get_skill("code-reproducer")
            if reproducer:
                reproduction = reproducer.apply(results["paper_info"], Context())
                results["reproduction"] = str(reproduction)
                _save_phase(workspace, "02_reproduction", reproduction)
                print(f"  ✅ [2/4] 复现完成")
            else:
                results["reproduction"] = "[code-reproducer skill 未注册]"
                print(f"  ⚠️  [2/4] 复现跳过 (skill 未注册)")
        except Exception as e:
            results["reproduction"] = f"[复现失败: {e}]"
            print(f"  ⚠️  [2/4] 复现失败: {e}")
    else:
        results["reproduction"] = "[跳过]"
        print(f"  ⏭️  [2/4] 跳过复现")

    # ── Phase 3: 实验 ──
    if "experiment" not in skip_phases:
        print(f"  🧪 [3/4] 实验...")
        try:
            experimenter = _get_skill("experimenter")
            if experimenter:
                experiment = experimenter.apply(
                    json.dumps({"paper": results["paper_info"], "reproduction": results["reproduction"]},
                              ensure_ascii=False),
                    Context()
                )
                results["experiment"] = str(experiment)
                _save_phase(workspace, "03_experiment", experiment)
                print(f"  ✅ [3/4] 实验完成")
            else:
                results["experiment"] = "[experimenter skill 未注册]"
                print(f"  ⚠️  [3/4] 实验跳过 (skill 未注册)")
        except Exception as e:
            results["experiment"] = f"[实验失败: {e}]"
            print(f"  ⚠️  [3/4] 实验失败: {e}")
    else:
        results["experiment"] = "[跳过]"
        print(f"  ⏭️  [3/4] 跳过实验")

    # ── Phase 4: 记录 ──
    if "record" not in skip_phases:
        print(f"  📝 [4/4] 记录笔记...")
        try:
            notebook = _get_skill("lab-notebook")
            if notebook:
                record = notebook.apply(
                    json.dumps(results, ensure_ascii=False),
                    Context()
                )
                results["record"] = str(record)
                _save_phase(workspace, "04_notebook", record)
                print(f"  ✅ [4/4] 笔记完成")
            else:
                # fallback: 自己写笔记
                results["record"] = _fallback_write_notebook(workspace, results)
                print(f"  ✅ [4/4] 笔记完成 (fallback)")
        except Exception as e:
            results["record"] = f"[记录失败: {e}]"
            print(f"  ⚠️  [4/4] 记录失败: {e}")
    else:
        results["record"] = "[跳过]"
        print(f"  ⏭️  [4/4] 跳过记录")

    elapsed = time.time() - t0

    # 汇总
    summary = {
        "paper": paper,
        "workspace": workspace,
        "phases_completed": [k for k, v in results.items() if "[跳过]" not in str(v) and "[失败" not in str(v)],
        "elapsed_seconds": round(elapsed, 1),
        "results_preview": {k: str(v)[:200] for k, v in results.items()},
    }

    return json.dumps(summary, ensure_ascii=False, indent=2)


def _get_skill(name: str):
    """从 SkillRegistry 获取 skill，如不存在则尝试注册"""
    try:
        from lambdagent.skills import SkillRegistry
        registry = SkillRegistry()
        skill = registry.get(name)
        if skill is None:
            # 尝试注册 research skill pack
            from lambdagent.skillpacks.research import register_all
            register_all()
            skill = registry.get(name)
        return skill
    except ImportError:
        return None


def _save_phase(workspace: str, phase_name: str, content) -> None:
    """保存阶段结果到工作空间"""
    phase_dir = os.path.join(workspace, phase_name)
    os.makedirs(phase_dir, exist_ok=True)
    output_path = os.path.join(phase_dir, "output.json")
    try:
        with open(output_path, "w", encoding="utf-8") as f:
            if isinstance(content, str):
                try:
                    data = json.loads(content)
                    json.dump(data, f, ensure_ascii=False, indent=2)
                except json.JSONDecodeError:
                    f.write(content)
            else:
                json.dump(str(content), f, ensure_ascii=False, indent=2)
    except Exception:
        pass


def _fallback_read_paper(paper: str) -> str:
    """当 paper-reader skill 不可用时的 fallback"""
    return json.dumps({
        "title": paper,
        "status": "需要手动搜索",
        "hint": "请用 WebSearch 搜索论文，然后 WebFetch 获取内容",
    }, ensure_ascii=False)


def _fallback_write_notebook(workspace: str, results: dict) -> str:
    """当 lab-notebook skill 不可用时的 fallback"""
    notebook_path = os.path.join(workspace, "research_notes.md")
    content = f"""# 研究笔记

## 论文信息
{results.get('paper_info', '无')}

## 复现结果
{results.get('reproduction', '无')}

## 实验结果
{results.get('experiment', '无')}

---
*由 ResearchWorkflow 自动生成*
"""
    try:
        with open(notebook_path, "w", encoding="utf-8") as f:
            f.write(content)
        return f"笔记已保存到: {notebook_path}"
    except Exception as e:
        return f"[保存失败: {e}]"