lambdagent/agentexample/research67/run_pipeline.py

#!/usr/bin/env python3
"""
research67/run_pipeline.py — 科研论文生成管线

多轮迭代：每轮执行一个阶段，评审不通过则回退修改，直到论文完成。

用法:
    python3 agentexample/research67/run_pipeline.py --idea "你的研究想法"
    python3 agentexample/research67/run_pipeline.py --idea-file path/to/IDEA.md
    python3 agentexample/research67/run_pipeline.py --claude   # 强制 Claude Code
    python3 agentexample/research67/run_pipeline.py --ollama   # 用本地 Ollama
"""
from __future__ import annotations

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

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

from lambdagent.fromconfig import from_config
from lambdagent.core import Context


# ══════════════════════════════════════════════════════
# 管线阶段定义
# ══════════════════════════════════════════════════════

PHASES = [
    {
        "name": "文献调研",
        "prompt": (
            "第一步：文献调研。\n"
            "1. 用 WebSearch 搜索与该研究主题相关的论文（至少搜索 3 次不同关键词）\n"
            "2. 整理搜索结果，提取每篇论文的标题、作者、年份、核心方法、主要结论\n"
            "3. 用 WriteFile 将文献综述写入 {workspace}/01_literature_review.md\n"
            "4. 总结研究现状和 gap"
        ),
    },
    {
        "name": "实验方案",
        "prompt": (
            "第二步：制定实验方案。\n"
            "1. 先用 ReadFile 读取 {workspace}/01_literature_review.md 了解前期调研\n"
            "2. 基于研究想法和文献综述，设计实验方案\n"
            "3. 包含：数据集选择、baseline 对比、评价指标、实验步骤\n"
            "4. 用 WriteFile 写入 {workspace}/02_experiment_plan.md"
        ),
    },
    {
        "name": "实验执行",
        "prompt": (
            "第三步：执行实验。\n"
            "1. 先读取 {workspace}/02_experiment_plan.md 了解实验方案\n"
            "2. 用 WriteFile 编写 Python 实验代码到 {workspace}/code/\n"
            "3. 用 Bash 运行实验代码\n"
            "4. 将实验结果（数据/图表）保存到 {workspace}/results/\n"
            "5. 用 WriteFile 写实验报告到 {workspace}/03_experiment_results.md"
        ),
    },
    {
        "name": "论文撰写",
        "prompt": (
            "第四步：撰写论文。\n"
            "1. 读取前几个阶段的文件：01_literature_review.md, 02_experiment_plan.md, 03_experiment_results.md\n"
            "2. 撰写完整的学术论文，包含：\n"
            "   - Abstract（200字以内）\n"
            "   - 1. Introduction（研究背景、问题、贡献）\n"
            "   - 2. Related Work（基于文献调研）\n"
            "   - 3. Method（提出的方法详细描述）\n"
            "   - 4. Experiments（实验设置、结果、分析）\n"
            "   - 5. Conclusion（总结和未来工作）\n"
            "   - References\n"
            "3. 用 WriteFile 写入 {workspace}/04_paper.md"
        ),
    },
    {
        "name": "论文评审",
        "prompt": (
            "第五步：模拟评审。\n"
            "1. 读取 {workspace}/04_paper.md\n"
            "2. 以**该领域世界最顶级期刊/会议**的资深审稿人视角评审论文。\n"
            "   根据论文所属学科自动匹配对标标准，例如：\n"
            "   - 计算机: NeurIPS / ICML / CVPR / ACL / SIGMOD / OSDI / SOSP\n"
            "   - 物理: Nature Physics / Physical Review Letters\n"
            "   - 生物医学: Nature / Science / Cell / The Lancet\n"
            "   - 数学: Annals of Mathematics / Inventiones Mathematicae\n"
            "   - 经济学: AER / Econometrica / QJE\n"
            "   - 材料/化学: Nature Materials / JACS / Angewandte Chemie\n"
            "   - 综合/交叉: Nature / Science / PNAS\n"
            "3. 从以下维度评分（1-5分）：\n"
            "   - novelty（新颖性）：相比 state-of-the-art 的创新程度\n"
            "   - soundness（技术正确性）：方法论、实验设计、统计分析是否严谨\n"
            "   - clarity（表述清晰度）：逻辑结构、语言质量、图表规范\n"
            "   - significance（研究意义）：对领域的潜在影响\n"
            "   - reproducibility（可复现性）：数据/代码/实验描述是否充分\n"
            "4. 计算综合评分（加权平均），给出 acceptance_probability (0-1)\n"
            "5. 指出论文匹配的最佳目标期刊/会议（及其投稿截止日期如果知道）\n"
            "6. 列出具体的修改建议（至少 3 条，按优先级排序）\n"
            "7. 用 WriteFile 将评审结果写入 {workspace}/05_review.json，格式：\n"
            '   {{"target_venue":"期刊/会议名","novelty":X,"soundness":X,"clarity":X,'
            '"significance":X,"reproducibility":X,"acceptance_probability":X,'
            '"suggestions":["..."],"strengths":["..."],"weaknesses":["..."],"summary":"..."}}'
        ),
    },
]

REVISION_PROMPT = (
    "论文评审未通过（acceptance_probability < {threshold}）。\n"
    "评审意见：{suggestions}\n\n"
    "请根据评审意见修改论文：\n"
    "1. 读取 {workspace}/04_paper.md 和 {workspace}/05_review.json\n"
    "2. 针对每条修改建议逐一改进\n"
    "3. 用 WriteFile 将修改后的论文覆盖写入 {workspace}/04_paper.md\n"
    "4. 简要说明你做了哪些修改"
)


# ══════════════════════════════════════════════════════
# 管线执行器
# ══════════════════════════════════════════════════════

def run_pipeline(idea: str, max_rounds: int = 3, threshold: float = 0.5,
                 config_path: str = None):
    """执行科研管线，多轮迭代直到论文通过或达到最大轮次。"""

    config_path = config_path or str(Path(__file__).parent / "agent-config.yml")

    # 创建工作区
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    workspace = str(Path(__file__).parent / "workspace" / f"run_{timestamp}")
    os.makedirs(workspace, exist_ok=True)
    os.makedirs(f"{workspace}/code", exist_ok=True)
    os.makedirs(f"{workspace}/results", exist_ok=True)

    # 保存 IDEA
    with open(f"{workspace}/IDEA.md", "w") as f:
        f.write(idea)

    print(f"\n{'═' * 60}")
    print(f"  📚 research67 — 科研论文生成管线")
    print(f"  工作区: {workspace}")
    print(f"  最大轮次: {max_rounds}")
    print(f"  通过阈值: {threshold}")
    print(f"{'═' * 60}\n")

    for round_num in range(1, max_rounds + 1):
        print(f"\n{'─' * 40}")
        print(f"  🔄 第 {round_num} 轮")
        print(f"{'─' * 40}")

        # 每轮创建新的 agent 实例（新会话）
        term = from_config(config_path)
        ctx = Context()

        # 首次输入包含研究想法
        initial_context = (
            f"研究主题：\n{idea}\n\n"
            f"工作区路径: {workspace}\n"
            f"当前轮次: {round_num}/{max_rounds}\n"
        )

        # 执行每个阶段
        for phase in PHASES:
            phase_name = phase["name"]
            phase_prompt = phase["prompt"].format(workspace=workspace)

            print(f"\n  📌 阶段: {phase_name}")
            t0 = time.time()

            if phase == PHASES[0]:
                # 首阶段：带完整上下文
                result = term.apply(initial_context + "\n" + phase_prompt, ctx)
            else:
                # 后续阶段：agent 已有记忆
                result = term.apply(phase_prompt, ctx)

            elapsed = time.time() - t0
            print(f"  ✅ {phase_name} 完成 ({elapsed:.0f}s, {len(ctx.trace)} 步)")

        # 检查评审结果
        review_path = f"{workspace}/05_review.json"
        if os.path.exists(review_path):
            try:
                with open(review_path) as f:
                    review = json.loads(f.read())
                score = review.get("acceptance_probability", 0)
                print(f"\n  📊 评审结果: acceptance_probability = {score:.2f}")
                print(f"     novelty={review.get('novelty')}, soundness={review.get('soundness')}, "
                      f"clarity={review.get('clarity')}, significance={review.get('significance')}")

                if score >= threshold:
                    print(f"\n  🎉 论文通过！(score={score:.2f} ≥ {threshold})")
                    break

                if round_num < max_rounds:
                    suggestions = review.get("suggestions", [])
                    print(f"  ⚠️  未通过 (score={score:.2f} < {threshold})，开始修改...")

                    # 修改论文
                    revision_prompt = REVISION_PROMPT.format(
                        threshold=threshold,
                        suggestions=json.dumps(suggestions, ensure_ascii=False),
                        workspace=workspace,
                    )
                    result = term.apply(revision_prompt, ctx)
                    print(f"  ✅ 修改完成")
            except (json.JSONDecodeError, KeyError) as e:
                print(f"  ⚠️  评审结果解析失败: {e}")
        else:
            print(f"  ⚠️  未生成评审文件: {review_path}")

    # 最终输出
    paper_path = f"{workspace}/04_paper.md"
    print(f"\n{'═' * 60}")
    print(f"  📚 管线结束")
    print(f"  论文路径: {paper_path}")
    print(f"  工作区: {workspace}")
    if os.path.exists(paper_path):
        with open(paper_path) as f:
            content = f.read()
        print(f"  论文长度: {len(content)} 字符, {content.count(chr(10))} 行")
    print(f"{'═' * 60}")


# ══════════════════════════════════════════════════════
# CLI 入口
# ══════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="📚 research67 科研论文生成管线")
    parser.add_argument("--idea", type=str, help="研究想法（直接文本）")
    parser.add_argument("--idea-file", type=str, help="研究想法文件路径")
    parser.add_argument("--max-rounds", type=int, default=3, help="最大迭代轮次（默认 3）")
    parser.add_argument("--threshold", type=float, default=0.5, help="通过阈值（默认 0.5）")
    parser.add_argument("--config", type=str, default=None, help="agent-config.yml 路径")
    args = parser.parse_args()

    if args.idea_file:
        with open(args.idea_file) as f:
            idea = f.read()
    elif args.idea:
        idea = args.idea
    else:
        print("请提供研究想法：")
        print("  --idea '研究主题描述'")
        print("  --idea-file path/to/IDEA.md")
        return

    run_pipeline(
        idea=idea,
        max_rounds=args.max_rounds,
        threshold=args.threshold,
        config_path=args.config,
    )


if __name__ == "__main__":
    main()