ranranrunforit/deep_research

1from agents import Runner, trace, gen_trace_id2from search_agent import search_agent3from planner_agent import planner_agent, WebSearchItem, WebSearchPlan4from writer_agent import writer_agent, ReportData5from email_agent import email_agent6import asyncio7 8class ResearchManager:9 10    async def run(self, query: str):11        """ Run the deep research process, yielding the status updates and the final report"""12        trace_id = gen_trace_id()13        with trace("Research trace", trace_id=trace_id):14            print(f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}")15            yield f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}"16            print("Starting research...")17            search_plan = await self.plan_searches(query)18            yield "Searches planned, starting to search..."     19            search_results = await self.perform_searches(search_plan)20            yield "Searches complete, writing report..."21            report = await self.write_report(query, search_results)22            yield "Report written, sending email..."23            await self.send_email(report)24            yield "Email sent, research complete"25            yield report.markdown_report26        27 28    async def plan_searches(self, query: str) -> WebSearchPlan:29        """ Plan the searches to perform for the query """30        print("Planning searches...")31        result = await Runner.run(32            planner_agent,33            f"Query: {query}",34        )35        print(f"Will perform {len(result.final_output.searches)} searches")36        return result.final_output_as(WebSearchPlan)37 38    async def perform_searches(self, search_plan: WebSearchPlan) -> list[str]:39        """ Perform the searches to perform for the query """40        print("Searching...")41        num_completed = 042        tasks = [asyncio.create_task(self.search(item)) for item in search_plan.searches]43        results = []44        for task in asyncio.as_completed(tasks):45            result = await task46            if result is not None:47                results.append(result)48            num_completed += 149            print(f"Searching... {num_completed}/{len(tasks)} completed")50        print("Finished searching")51        return results52 53    async def search(self, item: WebSearchItem) -> str | None:54        """ Perform a search for the query """55        input = f"Search term: {item.query}\nReason for searching: {item.reason}"56        try:57            result = await Runner.run(58                search_agent,59                input,60            )61            return str(result.final_output)62        except Exception:63            return None64 65    async def write_report(self, query: str, search_results: list[str]) -> ReportData:66        """ Write the report for the query """67        print("Thinking about report...")68        input = f"Original query: {query}\nSummarized search results: {search_results}"69        result = await Runner.run(70            writer_agent,71            input,72        )73 74        print("Finished writing report")75        return result.final_output_as(ReportData)76    77    async def send_email(self, report: ReportData) -> None:78        print("Writing email...")79        result = await Runner.run(80            email_agent,81            report.markdown_report,82        )83        print("Email sent")84        return report