File size: 3,875 Bytes
e83829b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
from agents import Runner, trace, gen_trace_id
from search_agent import search_agent
from planner_agent import planner_agent, WebSearchItem, WebSearchPlan
from questions_generator_agent import questions_generator_agent
from writer_agent import writer_agent, ReportData
from email_agent import email_agent
import asyncio
from agents import function_tool

"""

    async def run(self, query: str):

        " Run the deep research process, yielding the status updates and the final report"

        trace_id = gen_trace_id()

        with trace("Research trace", trace_id=trace_id):

            print(f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}")

            yield f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}"

            print("Starting research...")

            search_plan_result = await self.plan_searches(query)

            search_plan = search_plan_result.searches

            user_email = search_plan_result.user_email

            yield "Searches planned, starting to search..."     

            search_results = await self.perform_searches(search_plan)

            yield "Searches complete, writing report..."

            report = await self.write_report(query, search_results)

            yield "Report written, sending email..."

            await self.send_email(report, user_email)

            yield "Email sent, research complete"

            yield report.markdown_report

"""

@function_tool
async def generate_questions(query: str) -> str:
    """ Generate 5 questions to answer for the query """
    print("Generating questions...")
    result = await Runner.run(
        questions_generator_agent,
        f"Query: {query}",
    )
    print(f"Generated 5 questions...")
    return result.final_output

@function_tool
async def plan_searches(query: str) -> WebSearchPlan:
    """ Plan the searches to perform for the query """
    print("Planning searches...")
    result = await Runner.run(
        planner_agent,
        f"Query: {query}",
    )
    print(f"Will perform {len(result.final_output.searches)} searches")
    return result.final_output_as(WebSearchPlan)

@function_tool
async def perform_searches(search_plan: WebSearchPlan) -> list[str]:
    """ Perform the searches to perform for the query """
    print("Searching...")
    num_completed = 0
    tasks = [asyncio.create_task(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_completed += 1
        print(f"Searching... {num_completed}/{len(tasks)} completed")
    print("Finished searching")
    return results

async def search(item: WebSearchItem) -> str | None:
    """ Perform a search for the query """
    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

@function_tool
async def write_report(query: str, search_results: list[str]) -> ReportData:
    """ Write the report for the query """
    print("Thinking about report...")
    input = f"Original query: {query}\nSummarized search results: {search_results}"
    result = await Runner.run(
        writer_agent,
        input,
    )

    print("Finished writing report")
    return result.final_output_as(ReportData)
    
@function_tool
async def send_email(report: ReportData, user_email: str) -> None:
    print("Writing email...")
    result = await Runner.run(
        email_agent,
        f"Report: {report.markdown_report}\nUser email: {user_email}",
    )
    print("Email sent")
    return report