sagarnildass/Deep_Research_Assistant_Agent

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 asyncio7from typing import Optional8 9class ResearchManagerAgent:10 11    async def run(12        self,13        query: str,14        clarifying_questions: list[str],15        clarifying_answers: list[str],16        send_email_flag: bool = False,17        recipient_email: Optional[str] = None,18    ):19        """ Run the deep research process using user-provided clarification answers. """20        trace_id = gen_trace_id()21        with trace("Research trace", trace_id=trace_id):22            print(f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}")23            yield f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}"24            yield "Planning search based on clarifications..."25 26            print(f"Clarifying questions: {clarifying_questions}")27            print(f"Clarifying answers: {clarifying_answers}")28 29            # Plan searches using clarifications and user answers30            search_plan = await self.plan_searches(query, clarifying_questions, clarifying_answers)31 32            yield "Searches planned, starting to search..."33            search_results = await self.perform_searches(search_plan)34 35            yield "Searches complete, writing report..."36            report = await self.write_report(query, search_results)37 38            if send_email_flag and recipient_email:39                yield f"Sending report to {recipient_email}..."40                await self.send_email(report, recipient_email)41                yield "Email sent"42            else:43                yield "Skipping email step"44               45            yield "Email sent"46            yield report.markdown_report47 48    async def plan_searches(self, query: str, questions: list[str], answers: list[str]) -> WebSearchPlan:49        """ Plan the searches to perform based on clarifications """50        print("Planning searches...")51 52        # Combine clarifying Q&A into structured prompt53        clarifying_context = "\n".join(54            f"Q: {q}\nA: {a}" for q, a in zip(questions, answers)55        )56        final_prompt = f"Query: {query}\nClarifications:\n{clarifying_context}"57 58        result = await Runner.run(59            planner_agent,60            input=final_prompt,61        )62        print(f"Will perform {len(result.final_output.searches)} searches")63        return result.final_output_as(WebSearchPlan)64 65    async def perform_searches(self, search_plan: WebSearchPlan) -> list[str]:66        """ Perform the searches for the planned queries """67        print("Searching...")68        num_completed = 069        tasks = [asyncio.create_task(self.search(item)) for item in search_plan.searches]70        results = []71        for task in asyncio.as_completed(tasks):72            result = await task73            if result is not None:74                results.append(result)75            num_completed += 176            print(f"Searching... {num_completed}/{len(tasks)} completed")77        print("Finished searching")78        return results79 80    async def search(self, item: WebSearchItem) -> Optional[str]:81        """ Perform a single web search """82        input_text = f"Search term: {item.query}\nReason for searching: {item.reason}"83        try:84            result = await Runner.run(85                search_agent,86                input_text,87            )88            return str(result.final_output)89        except Exception as e:90            print(f"Search failed: {e}")91            return None92 93    async def write_report(self, query: str, search_results: list[str]) -> ReportData:94        """ Write a markdown report from search results """95        print("Thinking about report...")96        input_text = f"Original query: {query}\nSummarized search results: {search_results}"97        result = await Runner.run(98            writer_agent,99            input_text,100        )101        print("Finished writing report")102        return result.final_output_as(ReportData)103 104    async def send_email(self, report: ReportData, recipient_email: str) -> None:105        """ Send the report via email """106 107        email_prompt = f"""Send the following report as an email.108        To: {recipient_email}109        Body (HTML):110        {report.markdown_report}111        """112        print(f"Sending email to: {recipient_email}")113        await Runner.run(email_agent, input=email_prompt)114        print("✅ Email sent")