chore: import upstream snapshot with attribution

2026-07-13 12:39:17 +08:00
commit 4ed4e9ff99
1368 changed files with 334957 additions and 0 deletions
@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import asyncio
+import time
+
+from rich.console import Console
+
+from agents import Runner, custom_span, gen_trace_id, trace
+
+from .agents.planner_agent import WebSearchItem, WebSearchPlan, planner_agent
+from .agents.search_agent import search_agent
+from .agents.writer_agent import ReportData, writer_agent
+from .printer import Printer
+
+
+class ResearchManager:
+    def __init__(self):
+        self.console = Console()
+        self.printer = Printer(self.console)
+
+    async def run(self, query: str) -> None:
+        trace_id = gen_trace_id()
+        with trace("Research trace", trace_id=trace_id):
+            self.printer.update_item(
+                "trace_id",
+                f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}",
+                is_done=True,
+                hide_checkmark=True,
+            )
+
+            self.printer.update_item(
+                "starting",
+                "Starting research...",
+                is_done=True,
+                hide_checkmark=True,
+            )
+            search_plan = await self._plan_searches(query)
+            search_results = await self._perform_searches(search_plan)
+            report = await self._write_report(query, search_results)
+
+            final_report = f"Report summary\n\n{report.short_summary}"
+            self.printer.update_item("final_report", final_report, is_done=True)
+
+            self.printer.end()
+
+        print("\n\n=====REPORT=====\n\n")
+        print(f"Report: {report.markdown_report}")
+        print("\n\n=====FOLLOW UP QUESTIONS=====\n\n")
+        follow_up_questions = "\n".join(report.follow_up_questions)
+        print(f"Follow up questions: {follow_up_questions}")
+
+    async def _plan_searches(self, query: str) -> WebSearchPlan:
+        self.printer.update_item("planning", "Planning searches...")
+        result = await Runner.run(
+            planner_agent,
+            f"Query: {query}",
+        )
+        self.printer.update_item(
+            "planning",
+            f"Will perform {len(result.final_output.searches)} searches",
+            is_done=True,
+        )
+        return result.final_output_as(WebSearchPlan)
+
+    async def _perform_searches(self, search_plan: WebSearchPlan) -> list[str]:
+        with custom_span("Search the web"):
+            self.printer.update_item("searching", "Searching...")
+            num_completed = 0
+            num_succeeded = 0
+            num_failed = 0
+            tasks = [asyncio.create_task(self._search(item)) for item in search_plan.searches]
+            results = []
+            for task in asyncio.as_completed(tasks):
+                result = await task
+                if result is not None:
+                    results.append(result)
+                    num_succeeded += 1
+                else:
+                    num_failed += 1
+                num_completed += 1
+                status = f"Searching... {num_completed}/{len(tasks)} finished"
+                if num_failed:
+                    status += f" ({num_succeeded} succeeded, {num_failed} failed)"
+                self.printer.update_item(
+                    "searching",
+                    status,
+                )
+            summary = f"Searches finished: {num_succeeded}/{len(tasks)} succeeded"
+            if num_failed:
+                summary += f", {num_failed} failed"
+            self.printer.update_item("searching", summary, is_done=True)
+            return results
+
+    async def _search(self, item: WebSearchItem) -> str | None:
+        input = f"Search term: {item.query}\nReason for searching: {item.reason}"
+        try:
+            result = await Runner.run(
+                search_agent,
+                input,
+            )
+            return str(result.final_output)
+        except Exception:
+            return None
+
+    async def _write_report(self, query: str, search_results: list[str]) -> ReportData:
+        self.printer.update_item("writing", "Thinking about report...")
+        input = f"Original query: {query}\nSummarized search results: {search_results}"
+        result = Runner.run_streamed(
+            writer_agent,
+            input,
+        )
+        update_messages = [
+            "Thinking about report...",
+            "Planning report structure...",
+            "Writing outline...",
+            "Creating sections...",
+            "Cleaning up formatting...",
+            "Finalizing report...",
+            "Finishing report...",
+        ]
+
+        last_update = time.time()
+        next_message = 0
+        async for _ in result.stream_events():
+            if time.time() - last_update > 5 and next_message < len(update_messages):
+                self.printer.update_item("writing", update_messages[next_message])
+                next_message += 1
+                last_update = time.time()
+
+        self.printer.mark_item_done("writing")
+        return result.final_output_as(ReportData)