【Multi-agent实战】LangGraph 实现可视化的科研 Multi-Agent实战项目

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

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,
    )

# nodes.py
from typing import Dict
from langgraph.graph import END
from langchain.utilities import ArxivAPIWrapper
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

from research_state import ResearchState
from llm_config import build_llm

llm = build_llm()
arxiv = ArxivAPIWrapper()

# ===== 1) PaperHunter Node =====
def paper_hunter_node(state: ResearchState) -> ResearchState:
    topic = state["topic"]
    
    query = (
        f"{topic}, focus on works since 2022, "
        "especially those using CLIP or vision-language models."
    )
    result = arxiv.run(query)
    
    state["papers_raw"].append(result)
    state["logs"].append(f"[PaperHunter] searched arxiv with query: {query}")
    return state

# nodes.py (继续)
paper_analyze_prompt = PromptTemplate(
    input_variables=["paper_info"],
    template=(
        "你是一名科研助理，下面是若干论文的检索结果：\n{paper_info}\n\n"
        "请将这些论文整理为 JSON 数组，每个元素包含字段："
        'name, task, method, datasets, metrics, key_ideas, pros, cons。\n'
        "只输出 JSON，不要额外解释。"
    ),
)
paper_analyze_chain = LLMChain(llm=llm, prompt=paper_analyze_prompt)

def paper_analyst_node(state: ResearchState) -> ResearchState:
    # 简单地把 papers_raw 拼起来丢给总结链
    paper_info = "\n\n".join(state["papers_raw"])
    summary_json = paper_analyze_chain.run(paper_info=paper_info)
    
    state["papers_struct"].append(summary_json)
    state["logs"].append("[PaperAnalyst] summarized papers into structured JSON.")
    return state

# nodes.py (继续)
from langchain.prompts import PromptTemplate

experiment_prompt = PromptTemplate(
    input_variables=["paper_summary_json"],
    template=(
        "下面是关于某个研究主题的论文结构化总结（JSON）：\n{paper_summary_json}\n\n"
        "我的资源：单卡 RTX 3090（24GB），两周时间，每天训练时间约 6 小时。\n"
        "请设计一个可行的 baseline 实验方案，用 Markdown 输出，包含：\n"
        "1. 推荐复现的论文（1~2 篇）\n"
        "2. 数据集与预处理策略\n"
        "3. 模型与关键超参数\n"
        "4. 预估训练时间（粗略即可）\n"
        "5. TODO checklist（用 [ ] 列出来）\n"
    ),
)
experiment_chain = LLMChain(llm=llm, prompt=experiment_prompt)

def experiment_planner_node(state: ResearchState) -> ResearchState:
    # 简单起见，我们只拿最后一次的结构化结果
    paper_summary_json = state["papers_struct"][-1] if state["papers_struct"] else ""
    plan_md = experiment_chain.run(paper_summary_json=paper_summary_json)
    
    state["plan"] = plan_md
    state["logs"].append("[ExperimentPlanner] generated baseline experiment plan.")
    return state

def human_review_node(state: ResearchState) -> ResearchState:
    # 实战中可以把 state 持久化，等待前端/人工修改后再继续
    state["logs"].append("[HumanReview] checkpoint reached, please review papers_struct.")
    return state

# graph_build.py
from langgraph.graph import StateGraph, START, END

from research_state import ResearchState
from nodes import (
    paper_hunter_node,
    paper_analyst_node,
    experiment_planner_node,
    human_review_node,
)

def build_research_graph():
    builder = StateGraph(ResearchState)

    # 注册节点
    builder.add_node("paper_hunter", paper_hunter_node)
    builder.add_node("paper_analyst", paper_analyst_node)
    builder.add_node("human_review", human_review_node)
    builder.add_node("experiment_planner", experiment_planner_node)

    # 定义流程：START -> hunter -> analyst -> review -> planner -> END
    builder.add_edge(START, "paper_hunter")
    builder.add_edge("paper_hunter", "paper_analyst")
    builder.add_edge("paper_analyst", "human_review")
    builder.add_edge("human_review", "experiment_planner")
    builder.add_edge("experiment_planner", END)

    graph = builder.compile()
    return graph

# main.py
from graph_build import build_research_graph

def main():
    graph = build_research_graph()
    
    init_state = {
        "topic": "CLIP for deepfake detection and multimodal forgery since 2022",
        "papers_raw": [],
        "papers_struct": [],
        "plan": None,
        "logs": [],
    }

    final_state = graph.invoke(init_state)
    
    print("==== Logs ====")
    for log in final_state["logs"]:
        print(log)
    
    print("\n==== Plan ====\n")
    print(final_state["plan"])

if __name__ == "__main__":
    main()