Spaces:
Runtime error
Runtime error
File size: 4,737 Bytes
aa9134d |
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 |
from agents import Runner, trace, gen_trace_id
from model import model
from question_agent import question_agent, QuestionItem, QuestionPlan
from search_agent import search_agent
from planner_agent import planner_agent, WebSearchItem, WebSearchPlan
from writer_agent import writer_agent, ReportData
from refiner_agent import refiner_agent, ValidPlan
import asyncio
class ResearchManager:
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):
yield "Starting research"
question_plan: QuestionPlan = await self.ask_questions(query)
user_answers = {}
for q_item in question_plan.questions:
answer = input(f"{q_item.number}. {q_item.question}: ")
user_answers[q_item.number] = answer
q_item.answer = answer
yield f"Collected answers for {len(user_answers)} questions."
search_plan: WebSearchPlan = await self.plan_next(query, question_plan)
yield "Initial web search plan generated."
valid_plan: ValidPlan = await self.refine_plan(query, search_plan)
plan_num = 0
while not valid_plan.is_valid:
yield f"Plan not valid: {valid_plan.feedback}. Refining..."
search_plan: WebSearchPlan = await self.plan_next(query, question_plan)
valid_plan: ValidPlan = await self.refine_plan(query, search_plan)
if plan_num >= 4:
break
plan_num += 1
yield "Plan validated and refined."
search_results = await self.web_search(search_plan)
yield "Web searches completed."
report: ReportData = await self.write_report(query, search_results)
yield report.markdown_report
async def ask_questions(self, query: str)->QuestionPlan:
"""Given a user query, generate a set of 3 insightful questions to ask the user in order to facilitate detailed and in-depth planning."""
result = await Runner.run(question_agent, query)
return result.final_output_as(QuestionPlan)
async def plan_next(self, query:str, que_depth: QuestionPlan)->WebSearchPlan:
"""Based on the user query and the questions they answered, you come up with a set of web searches to perform to best answer the query"""
answers_str = "; ".join([q.answer for q in que_depth.questions])
input_query = f'User query {query}. User answered questions {answers_str}'
result = await Runner.run(planner_agent, input_query)
return result.final_output_as(WebSearchPlan)
async def refine_plan(self, query:str, ex_plan: WebSearchPlan)->ValidPlan:
"""Evaluate the research plan based on the user's query. """
all_plans = "; ".join([p.reason for p in ex_plan.searches])
input_plan = f'Validate the plan and provide whether its valid or not along with feedback and a bool of whether its valid or not. The plan {all_plans}'
result = await Runner.run(refiner_agent, input_plan)
return result.final_output_as(ValidPlan)
async def web_search(self, search_plan: WebSearchPlan)->list[str]:
""" Perform the searches to perform for the query """
num_completed = 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_completed += 1
print(f"Searching... {num_completed}/{len(tasks)} completed")
print("Finished searching")
return results
async def search(self, 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
async def write_report(self, 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)
|