【Multi-agent】智能科研助手实战项目

from config import api_type, api_key, api_base, api_version, model_name
from langchain.chat_models import AzureChatOpenAI

# 统一的 LLM，所有 Agent 共用
llm = AzureChatOpenAI(
    openai_api_base=api_base,
    openai_api_version=api_version,
    deployment_name=model_name,
    openai_api_key=api_key,
    openai_api_type=api_type,
    temperature=0.2,  # 科研场景偏保守一点
)

from langchain.agents import Tool
from langchain.tools.python.tool import PythonREPLTool
from langchain.python import PythonREPL
from langchain.utilities import ArxivAPIWrapper

# 1）arXiv 检索工具
arxiv = ArxivAPIWrapper()

arxiv_tool = Tool(
    name="arxiv_search",
    func=arxiv.run,
    description=(
        "用于根据科研主题检索最新 arXiv 论文。"
        "输入一个英文 query，例如 'CLIP for deepfake detection 2023', "
        "返回论文标题、作者和摘要等信息。"
    ),
)

# 2）Python REPL：方便做一些算力估算、参数计算
python_repl = PythonREPL()
python_tool = PythonREPLTool(python_repl=python_repl)

from langchain.agents import initialize_agent, AgentType

# PaperHunterAgent：负责检索相关论文
paper_hunter_agent = initialize_agent(
    tools=[arxiv_tool],
    llm=llm,
    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
)

# 为了让 Supervisor 调它，把它包成一个 Tool
paper_hunter_tool = Tool(
    name="paper_hunter_agent",
    func=paper_hunter_agent.run,
    description=(
        "文献猎手 Agent。负责根据给定研究主题检索相关 arXiv 论文，"
        "输出论文列表（标题 + 简介）。"
    ),
)

query = (
    "Find recent papers (since 2022) using CLIP or vision-language models "
    "for deepfake detection, AIGC safety, or multimodal forgery detection."
)
print(paper_hunter_agent(query))

from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

paper_analyze_prompt = PromptTemplate(
    input_variables=["paper_info"],
    template=(
        "你是一名计算机视觉领域的资深科研助理，当前在帮我做文献笔记。\n"
        "请阅读下面的论文信息，并按照如下 JSON 格式输出：\n"
        "{{\n"
        '  "task": "...",\n'
        '  "method": "...",\n'
        '  "datasets": ["..."],\n'
        '  "metrics": ["..."],\n'
        '  "key_ideas": ["..."],\n'
        '  "pros": ["..."],\n'
        '  "cons": ["..."]\n'
        "}}\n\n"
        "论文信息：\n{paper_info}\n"
    ),
)

paper_analyze_chain = LLMChain(
    llm=llm,
    prompt=paper_analyze_prompt,
)

def analyze_papers(paper_list_text: str) -> str:
    """
    输入：PaperHunterAgent 返回的一堆论文信息（字符串）
    输出：每篇论文的结构化 JSON 总结（整体仍然是字符串）
    """
    return paper_analyze_chain.run(paper_info=paper_list_text)

paper_analyst_tool = Tool(
    name="paper_analyst_agent",
    func=analyze_papers,
    description=(
        "论文分析 Agent。输入是多篇论文的文本信息，"
        "输出为结构化 JSON，总结任务、方法、数据集、指标和优缺点。"
    ),
)

from langchain.agents import AgentExecutor

experiment_prompt = PromptTemplate(
    input_variables=["paper_summary_json", "resource_info"],
    template=(
        "你现在是我的实验规划助理，帮我基于以下论文总结信息，设计一个可以在我资源条件下完成的 baseline 实验。\n\n"
        "【论文结构化信息】\n{paper_summary_json}\n\n"
        "【我的资源限制】\n{resource_info}\n\n"
        "请输出一个 Markdown 格式的实验规划，包含：\n"
        "1. 推荐复现的论文（1~2 篇）及理由\n"
        "2. 数据集选择与预处理策略\n"
        "3. 模型结构（可以简述）与关键超参数（batch size, lr, epochs 等）\n"
        "4. 预估训练时间（粗略即可）\n"
        "5. 实验 TODO checklist（用 [ ] 列出来）\n"
    ),
)

experiment_chain = LLMChain(
    llm=llm,
    prompt=experiment_prompt,
)

def plan_experiments(paper_summary_json: str) -> str:
    resource_info = (
        "单卡 RTX 3090，24GB 显存；可用时间约两周；"
        "日常还要写代码和写论文，因此每天训练时间控制在 6 小时以内。"
    )
    return experiment_chain.run(
        paper_summary_json=paper_summary_json,
        resource_info=resource_info,
    )

experiment_planner_tool = Tool(
    name="experiment_planner_agent",
    func=plan_experiments,
    description=(
        "实验规划 Agent。输入为论文的结构化 JSON 总结，"
        "输出 Markdown 格式的 baseline 实验设计方案和 TODO 列表。"
    ),
)

note_prompt = PromptTemplate(
    input_variables=["topic", "paper_summary_json", "experiment_plan"],
    template=(
        "你是一位科研人员，正在整理一份关于主题「{topic}」的科研备忘录。\n"
        "请根据下面的论文总结和实验规划，用 Markdown 输出一份条理清晰的笔记：\n\n"
        "【论文总结 JSON】\n{paper_summary_json}\n\n"
        "【实验规划】\n{experiment_plan}\n\n"
        "结构建议：\n"
        "1. 研究主题与问题背景\n"
        "2. 相关工作简要对比（按论文列点）\n"
        "3. 适合在我本地跑的 baseline 方案\n"
        "4. 后续可能的改进方向\n"
        "注意：风格稍微口语一点，就像在给未来的自己写备忘录。"
    ),
)

note_chain = LLMChain(
    llm=llm,
    prompt=note_prompt,
)

def write_research_note(topic: str, paper_summary_json: str, experiment_plan: str) -> str:
    return note_chain.run(
        topic=topic,
        paper_summary_json=paper_summary_json,
        experiment_plan=experiment_plan,
    )

note_writer_tool = Tool(
    name="note_writer_agent",
    func=lambda x: write_research_note(
        topic="CLIP for Deepfake Detection",
        paper_summary_json=x,
        experiment_plan="[实验规划占位符：在 Supervisor 中替换为真实内容]",
    ),
    description=(
        "科研笔记 Agent。输入论文总结 JSON，"
        "输出一份 Markdown 形式的科研备忘录草稿。"
    ),
)

from langchain.agents import initialize_agent, AgentType

supervisor_tools = [
    paper_hunter_tool,
    paper_analyst_tool,
    experiment_planner_tool,
    # note_writer_tool 也可以在最后用，这里先不挂，演示清晰一点
]

supervisor_system_msg = """
你是一个科研项目的总控 Agent，负责调用其他子 Agent 完成任务。

子 Agent 列表：
1. paper_hunter_agent：根据主题检索相关论文。
2. paper_analyst_agent：对论文列表做结构化总结（JSON）。
3. experiment_planner_agent：根据论文总结和资源限制，设计 baseline 实验方案。

原则：
- 步骤化执行：先查文献，再结构化总结，再做实验规划。
- 输出时，请给出三个部分：
  (A) 论文列表简要整理
  (B) JSON 形式的论文结构化总结（可以缩写）
  (C) 实验规划要点（Markdown）

如果你认为任务已经完成，请直接输出最终结果，不要再调用工具。
"""

research_supervisor = initialize_agent(
    tools=supervisor_tools,
    llm=llm,
    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    agent_kwargs={
        "system_message": supervisor_system_msg,
    },
)

topic = "CLIP for Deepfake / multimodal forgery detection since 2022"
result = research_supervisor.run(
    f"请围绕主题「{topic}」，帮我完成文献检索 + 论文结构化总结 + baseline 实验规划。"
)
print(result)

> Entering new AgentExecutor chain...
Thought: 先用 paper_hunter_agent 找相关论文
Action:
{
  "action": "paper_hunter_agent",
  "action_input": "CLIP for deepfake detection, AIGC safety, multimodal forgery, since 2022"
}
Observation: 
[返回了一堆 arXiv 论文标题 + 摘要]

Thought: 需要把这些论文交给 paper_analyst_agent 做结构化总结
Action:
{
  "action": "paper_analyst_agent",
  "action_input": "..."   # 上一步的结果
}
Observation:
{ "papers": [ { "task": "...", "method": "...", ... }, ... ] }

Thought: 现在基础信息已经足够，可以调用 experiment_planner_agent 设计 baseline
Action:
{
  "action": "experiment_planner_agent",
  "action_input": "{ \"papers\": [ ... ] }"
}
Observation:
# Markdown 格式的实验规划...

Thought: 我已经有论文列表、结构化总结和实验规划，可以输出最终答案
Final Answer:
(A) 论文列表...
(B) 结构化 JSON...
(C) baseline 实验规划...
> Finished chain.

(A) 相关论文列表（部分）

1. 《CLIP-Guard: Text-Guided Defense against Deepfake Images》 (ICLR 2024)
   - 利用 CLIP 将文本描述与图像伪造特征对齐，实现可解释的伪造检测。
2. 《VL-Detect: Vision-Language Pretraining for Generalizable Deepfake Detection》 (CVPR 2023)
   - 将伪造类型描述编码到文本空间，与图像特征做对比学习。
3. 《Prompting Vision-Language Models for Face Forgery Detection》 (arXiv 2023)
   - 探索 prompt 设计对 CLIP deepfake detection 性能的影响。

(B) 论文结构化总结（节选）

[
  {
    "name": "CLIP-Guard",
    "task": "image-level deepfake detection + explanation",
    "method": "use CLIP text prompts to highlight suspicious regions and train a detector on CLIP features",
    "datasets": ["FF++", "CelebDF-v2"],
    "metrics": ["AUC", "EER"],
    "key_ideas": [
      "align textual descriptions of forgery cues with visual features",
      "leverage CLIP as a frozen feature extractor"
    ],
    "pros": [
      "good cross-dataset generalization",
      "provides human-understandable explanations"
    ],
    "cons": [
      "relies on carefully crafted prompts",
      "limited support for video temporal cues"
    ]
  },
  ...
]

(C) baseline 实验规划（概要）

- 推荐复现的论文：
  - 首选：VL-Detect（模型结构较简单，便于在 3090 上快速复现）
  - 备选：Prompting VLMs for Face Forgery Detection（适合做 prompt ablation）

- 数据集选择：
  - 训练：FF++ (c23)，只用部分 manipulation 类型（Deepfakes, FaceSwap）
  - 测试：
    - in-domain: FF++ test split
    - cross-domain: CelebDF-v2

- 训练配置建议：
  - backbone: CLIP ViT-B/16, image encoder frozen
  - classifier: 2-layer MLP on top of CLIP embedding
  - batch size: 64 (根据 24GB 显存估算)
  - lr: 1e-4, optimizer: AdamW
  - epochs: 20（约两周内可以完成多轮试验）

- TODO checklist:
  - [ ] 写数据加载脚本（对齐 FF++ / CelebDF-v2 的人脸裁剪方式）
  - [ ] 实现 CLIP 特征抽取模块（可复用 open_clip/pytorch-lightning 脚手架）
  - [ ] 训练 baseline classifier，并在 in-domain / cross-domain 上评估 AUC/EER
  - [ ] 做一轮 prompt ablation：比较手工 prompt vs 学习式 prompt
  - [ ] 整理结果和可视化（t-SNE, ROC curve）

from config import api_type, api_key, api_base, api_version, model_name

from langchain.chat_models import AzureChatOpenAI
from langchain.agents import (
    initialize_agent,
    AgentType,
    Tool,
)
from langchain.tools.python.tool import PythonREPLTool
from langchain.python import PythonREPL
from langchain.utilities import ArxivAPIWrapper
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain


def build_llm():
    return AzureChatOpenAI(
        openai_api_base=api_base,
        openai_api_version=api_version,
        deployment_name=model_name,
        openai_api_key=api_key,
        openai_api_type=api_type,
        temperature=0.2,
    )


def build_base_tools():
    arxiv = ArxivAPIWrapper()
    arxiv_tool = Tool(
        name="arxiv_search",
        func=arxiv.run,
        description="根据科研主题检索 arXiv 论文，输入英文 query。",
    )

    python_repl = PythonREPL()
    python_tool = PythonREPLTool(python_repl=python_repl)

    return arxiv_tool, python_tool


def build_paper_hunter_agent(llm, arxiv_tool):
    agent = initialize_agent(
        tools=[arxiv_tool],
        llm=llm,
        agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
    )
    tool = Tool(
        name="paper_hunter_agent",
        func=agent.run,
        description="文献猎手 Agent，用于根据主题检索相关 arXiv 论文。",
    )
    return agent, tool


def build_paper_analyst_tool(llm):
    prompt = PromptTemplate(
        input_variables=["paper_info"],
        template=(
            "你是一名资深科研助理，请对下面的论文信息做结构化总结，"
            "以 JSON 数组形式输出，每个元素包含字段："
            'name, task, method, datasets, metrics, key_ideas, pros, cons。\n\n'
            "论文信息：\n{paper_info}\n"
        ),
    )
    chain = LLMChain(llm=llm, prompt=prompt)

    def analyze(paper_info: str) -> str:
        return chain.run(paper_info=paper_info)

    tool = Tool(
        name="paper_analyst_agent",
        func=analyze,
        description="论文分析 Agent，将论文文本信息转为结构化 JSON。",
    )
    return tool


def build_experiment_planner_tool(llm):
    prompt = PromptTemplate(
        input_variables=["paper_summary_json"],
        template=(
            "下面是若干论文的结构化总结（JSON）：\n{paper_summary_json}\n\n"
            "假设我只有单卡 RTX 3090（24GB 显存），两周时间。\n"
            "请为我设计一个可行的 baseline 实验方案，用 Markdown 输出："
            "包括推荐复现的论文、数据集选择、模型配置、超参数和 TODO 列表。"
        ),
    )
    chain = LLMChain(llm=llm, prompt=prompt)

    def plan(paper_summary_json: str) -> str:
        return chain.run(paper_summary_json=paper_summary_json)

    tool = Tool(
        name="experiment_planner_agent",
        func=plan,
        description="实验规划 Agent，根据论文 JSON 设计 baseline 实验方案。",
    )
    return tool


def build_supervisor(llm, paper_hunter_tool, paper_analyst_tool, experiment_planner_tool):
    tools = [paper_hunter_tool, paper_analyst_tool, experiment_planner_tool]

    system_msg = """
你是科研项目的总控 Agent，需要调用其他子 Agent 完成任务：
1. 先调用 paper_hunter_agent 检索论文；
2. 再调用 paper_analyst_agent 做结构化总结；
3. 最后调用 experiment_planner_agent 设计 baseline 实验方案。

最终输出请包含：
(A) 论文列表简要整理
(B) JSON 形式的论文结构化总结（可以缩写）
(C) baseline 实验规划要点
"""

    supervisor = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
        agent_kwargs={"system_message": system_msg},
    )
    return supervisor


def main():
    llm = build_llm()
    arxiv_tool, python_tool = build_base_tools()

    _, paper_hunter_tool = build_paper_hunter_agent(llm, arxiv_tool)
    paper_analyst_tool = build_paper_analyst_tool(llm)
    experiment_planner_tool = build_experiment_planner_tool(llm)

    supervisor = build_supervisor(
        llm,
        paper_hunter_tool,
        paper_analyst_tool,
        experiment_planner_tool,
    )

    topic = "CLIP for deepfake detection and multimodal forgery since 2022"
    query = (
        f"请围绕主题「{topic}」，完成："
        "文献检索 + 论文结构化总结 + baseline 实验规划。"
    )
    result = supervisor.run(query)
    print("\n===== 最终结果 =====\n")
    print(result)


if __name__ == "__main__":
    main()