让智能体
自己决定跑几轮

from pydantic import BaseModel, Field


class SubTopic(BaseModel):
    """研究计划里的一个子话题。"""
    name: str
    status: str = "pending"  # pending（待办）| in_progress（进行中）| complete（已完成）
    notes: str = ""


class ResearchPlan(BaseModel):
    """指导整个研究过程的战略计划。"""
    sub_topics: list[SubTopic] = Field(default_factory=list)
    strategy_notes: str = ""

    @property
    def progress_summary(self) -> str:
        if not self.sub_topics:
            return "还没有计划"
        total = len(self.sub_topics)
        complete = sum(1 for st in self.sub_topics
                       if st.status == "complete")
        in_progress = sum(1 for st in self.sub_topics
                          if st.status == "in_progress")
        return (f"{complete}/{total} 已完成，"
                f"{in_progress} 进行中")


class ResearchState(BaseModel):
    goal: str = ""
    findings: list[str] = Field(default_factory=list)
    sources_consulted: list[str] = Field(default_factory=list)
    follow_up_questions: list[str] = Field(default_factory=list)
    plan: ResearchPlan = Field(default_factory=ResearchPlan)
    iteration_count: int = 0
    max_iterations: int = 10
    status: str = "in_progress"

    @property
    def should_continue(self) -> bool:
        # ↓ 这就是整章的终止门：三个条件缺一不可
        return (
            self.status == "in_progress"
            and self.iteration_count < self.max_iterations
            and len(self.follow_up_questions) > 0
        )

class ResearchIteration(BaseModel):
    summary_of_findings: str      # 这一轮查到了什么
    sources_used: list[str] = []  # 用了哪些来源
    follow_up_questions: list[str] = []   # ← 下一轮查什么，也是终止门第三条
    goal_satisfied: bool = False  # ← 够了吗，这是模型自己按的停止键
    confidence: float = 0.0
    plan_updates: list[SubTopicUpdate] = []
    new_sub_topics: list[SubTopicUpdate] = []
    strategy_notes: str = ""

import asyncio
import os

from agents import Agent, Runner, set_default_openai_api, set_tracing_disabled
from agents.mcp import MCPServerStdio
from dotenv import load_dotenv

# 仓库原文没有这三行（这一章一个 load_dotenv 都没有），是为了接第三方模型加的
load_dotenv()
set_default_openai_api("chat_completions")
set_tracing_disabled(True)

# …… 中间是从 01 和 03 复制过来的模型定义，省略 ……

research_agent = Agent(
    name="Deep Research Agent",
    instructions="""
你是一个深度研究智能体。
你的目标是通过检索信息、阅读结果、综合发现，把一个话题研究透彻。

研究计划：
- 第一轮（还没有计划时），分析研究目标，通过 new_sub_topics 返回若干子话题
  来建立计划。每个子话题应该是目标的一个不同侧面。
  未开始的设为 "pending"，当前正在研究的设为 "in_progress"。
- 后续每一轮，通过 plan_updates 更新已有的子话题（把状态改成 "in_progress"
  或 "complete"，并在 notes 里总结学到了什么）。如果发现研究需要涉及计划外的
  领域，也可以通过 new_sub_topics 添加新的子话题。
- 用 strategy_notes 记录你对研究方法的高层观察，或者对策略的调整。

每一轮，回顾你的计划和已积累的状态，找出空白，检索新信息，更新你的发现。
follow_up_questions 里要放下一轮具体的检索查询。
计划提供战略方向，follow_up_questions 驱动战术执行。

当你认为已经有足够的信息来全面回答研究目标时，把 goal_satisfied 设为 true。
""",
    model=os.environ["OPENAI_DEFAULT_MODEL"],   # 书里写死的是 gpt-4o
    output_type=ResearchIteration,
)


async def run_research_loop(goal: str, max_iterations: int = 10):
    search_server = MCPServerStdio(
        name="brave-search",
        params={
            "command": "npx",
            # 仓库原文写的是 @anthropic/brave-search-mcp —— 这个包不存在！
            "args": ["-y", "@modelcontextprotocol/server-brave-search"],
            "env": {"BRAVE_API_KEY": os.environ["BRAVE_API_KEY"]},
        },
    )
    async with search_server:
        agent = research_agent.clone(mcp_servers=[search_server])
        state = ResearchState(
            goal=goal,
            max_iterations=max_iterations,
            follow_up_questions=[goal],   # ← 拿目标本身当第一个问题，把循环点着
        )

        while state.should_continue:
            state.iteration_count += 1
            question = state.follow_up_questions.pop(0)   # 从队头取一个
            input_context = dict(
                current_question=question,
                accumulated_state=state.to_context(),     # 把状态讲给模型听
                research_plan=state.plan.to_context(),
            )
            print(f"\n--- 第 {state.iteration_count} 轮 ---")
            print(f"正在调查：{question}")

            result = await Runner.run(agent, input=str(input_context))
            iteration = result.final_output

            # 把这一轮的产出并回状态
            state.findings.append(iteration.summary_of_findings)
            state.sources_consulted.extend(iteration.sources_used)
            state.follow_up_questions.extend(iteration.follow_up_questions)
            apply_plan_updates(state.plan, iteration)

            if iteration.goal_satisfied:
                state.status = "complete"       # ← 模型自己按下停止键
                print(f"目标已达成（置信度：{iteration.confidence:.0%}）")

            print(f"计划进度：{state.plan.progress_summary}")

        return state