基于LangGraph搭建自己的rebattle小组

from typing import TypedDict, List, Dict, Any
​
class ParsedComment(TypedDict):
    id: str
    reviewer: str
    raw_text: str
    type: str          # 'experiment' / 'theory' / 'writing' ...
    severity: str      # 'major' / 'minor'
    summary: str
​
class PlanItem(TypedDict):
    id: str            # 对应 comment_id
    strategy: str      # 'refute' / 'accept' / 'partial' / 'clarify'
    actions: List[str] # ['recompute ablation', 'add limitation section' ...]
    notes: str         # 给自己看的审稿策略说明
​
class DraftReply(TypedDict):
    id: str
    content: str
    tokens: int
​
class RebuttalState(TypedDict):
    paper_title: str
    paper_abstract: str
    paper_key_ideas: str
    raw_reviews: str              # 原始所有 reviewer 意见
    parsed_comments: List[ParsedComment]
    reply_plan: List[PlanItem]
    draft_replies: Dict[str, DraftReply]
    final_rebuttal: str
    token_budget: int             # rebuttal 总 token 限制（例如 4000 字符）

# agents/parser_agent.py
import json
from typing import cast, List
from langchain_openai import ChatOpenAI
from .state import RebuttalState, ParsedComment
​
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.0)
​
def parser_node(state: RebuttalState) -> RebuttalState:
    """解析 raw_reviews -> parsed_comments"""
    prompt = f"""
你是一个学术审稿意见解析助手。
​
当前论文标题：{state['paper_title']}
​
以下是所有审稿意见（包含多个 reviewer）：
========
{state['raw_reviews']}
========
​
请你：
1. 按 reviewer 拆分意见（例如 R1 / R2 / R3）；
2. 按 bullet 或自然段将意见拆分为若干条 comment；
3. 对每条 comment 标注：
   - id: "R1-1", "R1-2" 这种
   - reviewer: "R1" / "R2" / ...
   - type: 从 ["experiment", "theory", "writing", "related_work", "clarification", "other"] 中选
   - severity: "major" 或 "minor"
   - summary: 用一句话中文概括该意见
​
请以 JSON 列表输出，比如：
[
  {{
    "id": "R1-1",
    "reviewer": "R1",
    "type": "experiment",
    "severity": "major",
    "summary": "认为跨数据集实验不充分，需要在 CelebDF-v2 上补充结果",
    "raw_text": "..."
  }},
  ...
]
只输出 JSON，不要加多余文字。
"""
    resp = llm.invoke(prompt)
    try:
        comments = json.loads(resp.content)
    except Exception:
        comments = []
​
    new_state = dict(state)
    new_state["parsed_comments"] = cast(List[ParsedComment], comments)
    return cast(RebuttalState, new_state)

# agents/planner_agent.py
from typing import List, cast
from langchain_openai import ChatOpenAI
from .state import RebuttalState, PlanItem, ParsedComment
import json
​
planner_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)
​
def planner_node(state: RebuttalState) -> RebuttalState:
    """根据 parsed_comments 生成 reply_plan"""
    comments: List[ParsedComment] = state["parsed_comments"]
​
    prompt = f"""
你是一个论文 Rebuttal 策略助手，擅长在保证礼貌的前提下，尽量维护论文评分。
​
下面是已经解析好的审稿 comment 列表（JSON）：
{json.dumps(comments, ensure_ascii=False, indent=2)}
​
请针对每条 comment 生成一个策略 PlanItem：
- id: 原样照抄 comment.id
- strategy: 从 ["refute", "accept", "partial", "clarify"] 中选
    - refute: 认为审稿人误解或说错了，要据理力争
    - accept: 明确承认问题，并说明已在修改中解决
    - partial: 部分同意，解释限制条件或 trade-off
    - clarify: 主要是澄清表述或补充说明
- actions: 列表，列出你建议的具体动作（用中文），例如：
    - "在附录中补充 CelebDF-v2 的 cross-dataset 实验"
    - "在方法部分补充公式推导细节"
- notes: 写给作者看的“策略说明”，比如
    - "该意见误读了我们的设置，可以通过强调 XXX 来澄清"
​
只输出一个 JSON 列表。
"""
    resp = planner_llm.invoke(prompt)
    try:
        plans = json.loads(resp.content)
    except Exception:
        plans = []
​
    new_state = dict(state)
    new_state["reply_plan"] = cast(List[PlanItem], plans)
    return cast(RebuttalState, new_state)

# agents/writer_agent.py
from typing import Dict, List, cast
from langchain_openai import ChatOpenAI
from .state import RebuttalState, ParsedComment, PlanItem, DraftReply
import json
​
writer_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)
​
def writer_node(state: RebuttalState) -> RebuttalState:
    """根据 reply_plan 逐条生成 draft_replies"""
    comments: List[ParsedComment] = state["parsed_comments"]
    plans: List[PlanItem] = state["reply_plan"]
    token_budget = state.get("token_budget", 4000)
​
    # 建立 id -> comment / plan 索引
    comment_map = {c["id"]: c for c in comments}
    plan_map = {p["id"]: p for p in plans}
​
    draft_replies: Dict[str, DraftReply] = {}
​
    for cid, comment in comment_map.items():
        plan = plan_map.get(cid)
        if plan is None:
            continue
​
        prompt = f"""
你现在在写论文 Rebuttal 的单条回复段落。
​
论文标题：{state['paper_title']}
论文关键想法：{state['paper_key_ideas']}
​
审稿意见（comment）：
{json.dumps(comment, ensure_ascii=False, indent=2)}
​
处理策略（plan）：
{json.dumps(plan, ensure_ascii=False, indent=2)}
​
请用英文写一段针对该 comment 的回复，要求：
- 开头先简要感谢 reviewer
- 根据 strategy:
  - refute: 礼貌但明确地指出 reviewer 的误解，并给出证据
  - accept: 承认问题并说明已经或将如何修改
  - partial: 部分同意，解释为什么在当前限制下这么做是合理的
  - clarify: 主要澄清文字或设置，保持语气温和
- 不要超过 220 词
- 语气专业、礼貌、有逻辑
​
直接输出正文，不要附加解释。
"""
        resp = writer_llm.invoke(prompt)
        content = resp.content
​
        draft_replies[cid] = {
            "id": cid,
            "content": content,
            "tokens": len(content.split()),
        }
​
    new_state = dict(state)
    new_state["draft_replies"] = cast(Dict[str, DraftReply], draft_replies)
    return cast(RebuttalState, new_state)

# agents/polisher_agent.py
from typing import Dict, List, cast
from langchain_openai import ChatOpenAI
from .state import RebuttalState, DraftReply, ParsedComment
import json
​
polish_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)
​
def polisher_node(state: RebuttalState) -> RebuttalState:
    comments: List[ParsedComment] = state["parsed_comments"]
    replies: Dict[str, DraftReply] = state["draft_replies"]
    budget = state.get("token_budget", 4000)
​
    # 按 reviewer 分组
    grouped: Dict[str, List[Dict]] = {}
    for c in comments:
        rid = c["reviewer"]
        grouped.setdefault(rid, [])
        rep = replies.get(c["id"])
        if rep is None:
            continue
        grouped[rid].append({
            "comment_id": c["id"],
            "comment_summary": c["summary"],
            "reply": rep["content"],
        })
​
    prompt = f"""
你现在要将下面的逐条 Rebuttal 回复整合成一份正式的 rebuttal 文档。
​
要求：
- 按 reviewer 分块（Reviewer #1, #2, #3）
- 对每个 comment:
  - 先引用一小段 reviewer 的 summary
  - 再给出我们的回复
- 控制整体长度不要太长（假设总限制约为 {budget} 词）
- 语气统一、礼貌、专业
- 可以适度合并相似的 comment，避免重复
​
下面是按 reviewer 分组的所有草稿回复（JSON）：
{json.dumps(grouped, ensure_ascii=False, indent=2)}
​
请直接输出最终 rebuttal 正文（英文）。
"""
    resp = polish_llm.invoke(prompt)
    rebuttal = resp.content
​
    new_state = dict(state)
    new_state["final_rebuttal"] = rebuttal
    return cast(RebuttalState, new_state)

# graph/build_rebuttal_graph.py
from langgraph.graph import StateGraph, START, END
from .state import RebuttalState
from agents.parser_agent import parser_node
from agents.planner_agent import planner_node
from agents.writer_agent import writer_node
from agents.polisher_agent import polisher_node
​
def build_rebuttal_graph():
    g = StateGraph(RebuttalState)
​
    g.add_node("parser", parser_node)
    g.add_node("planner", planner_node)
    g.add_node("writer", writer_node)
    g.add_node("polisher", polisher_node)
​
    g.add_edge(START, "parser")
    g.add_edge("parser", "planner")
    g.add_edge("planner", "writer")
    g.add_edge("writer", "polisher")
    g.add_edge("polisher", END)
​
    return g.compile()

# run_rebuttal.py
from graph.build_rebuttal_graph import build_rebuttal_graph
from state import RebuttalState
​
if __name__ == "__main__":
    # 这里的 raw_reviews / paper_xxx 都可以从本地文件读取
    raw_reviews = open("reviews.txt", "r", encoding="utf-8").read()
    paper_abstract = open("abstract.txt", "r", encoding="utf-8").read()
    key_ideas = "我们提出 FFE 模块从潜空间中显式提取 forgery factor，并基于能量分区做 targeted augmentation..."
​
    init_state: RebuttalState = {
        "paper_title": "Energy-aware Forgery Factor Extraction for Generalizable Deepfake Detection",
        "paper_abstract": paper_abstract,
        "paper_key_ideas": key_ideas,
        "raw_reviews": raw_reviews,
        "parsed_comments": [],
        "reply_plan": [],
        "draft_replies": {},
        "final_rebuttal": "",
        "token_budget": 3800,
    }
​
    graph = build_rebuttal_graph()
    final_state = graph.invoke(init_state)
​
    print("==== Final Rebuttal ====")
    print(final_state["final_rebuttal"])
​
    with open("rebuttal_draft.md", "w", encoding="utf-8") as f:
        f.write(final_state["final_rebuttal"])