jmc310/deep_research

1from agents import Runner, trace, gen_trace_id
2from search_agent import search_agent
3from planner_agent import planner_agent, WebSearchItem, WebSearchPlan
4from writer_agent import writer_agent, ReportData
5from email_agent import email_agent
6import asyncio
7
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_report
26        
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 = 0
42        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 task
46            if result is not None:
47                results.append(result)
48            num_completed += 1
49            print(f"Searching... {num_completed}/{len(tasks)} completed")
50        print("Finished searching")
51        return results
52
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 None
64
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