from agents import Runner, trace, gen_trace_id, Agent, function_tool
from search_agent import search_agent, create_search_agent
from planner_agent import planner_agent, create_planner_agent, WebSearchItem, WebSearchPlan
from writer_agent import writer_agent, create_writer_agent, ReportData
from email_agent import email_agent, create_email_agent
from clarifier_agent import clarifier_agent, create_clarifier_agent, ClarificationData
from evaluator_agent import evaluator_agent, optimizer_agent, create_evaluator_agent, create_optimizer_agent, EvaluationResult, OptimizedReport
import asyncio
from typing import Dict, Any, AsyncGenerator
# Legacy ResearchManager class for backward compatibility
class ResearchManager:
def __init__(self, api_key: str = None, model: str = "gpt-4o-mini"):
"""Initialize ResearchManager with optional API key and model"""
self.api_key = api_key
self.model = model
# Set the API key in environment if provided
if api_key:
import os
os.environ["OPENAI_API_KEY"] = api_key
async def run_with_clarification(self, query: str):
""" Run the clarification step and return clarifying questions """
trace_id = gen_trace_id()
with trace("Clarification trace", trace_id=trace_id):
print(f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}")
print("Generating clarifying questions...")
clarifier = create_clarifier_agent(self.model)
result = await Runner.run(
clarifier,
f"Query: {query}",
)
clarification_data = result.final_output_as(ClarificationData)
print(f"Generated {len(clarification_data.questions)} clarifying questions")
return {
"questions": clarification_data.questions,
"trace_id": trace_id
}
async def run_research_with_answers(self, query: str, answers: list[str]):
""" Run the full research process with clarification answers """
trace_id = gen_trace_id()
with trace("Research with clarification trace", trace_id=trace_id):
print(f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}")
print("Starting research with clarifications...")
# Use the new manager agent instead
clarified_query = self._format_clarified_query(query, answers)
result = await Runner.run(
ResearchManagerAgent,
f"Research Query: {clarified_query}",
)
return {
"report": result.final_output,
"trace_id": trace_id
}
def _format_clarified_query(self, original_query: str, answers: list[str]) -> str:
""" Format the original query with clarification answers """
clarifications = []
for i, answer in enumerate(answers, 1):
if answer.strip():
clarifications.append(f"{i}. {answer.strip()}")
if clarifications:
clarified_query = f"""Original query: {original_query}
Clarifications provided:
{chr(10).join(clarifications)}
Please use these clarifications to focus and refine the research approach."""
else:
clarified_query = original_query
return clarified_query
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...")
# Use the new manager agent
result = await Runner.run(ResearchManagerAgent, f"Research Query: {query}")
yield "Research complete"
yield result.final_output
# Global variable to store current model
_current_model = "gpt-4o-mini"
def set_current_model(model: str):
"""Set the current model for function tools"""
global _current_model
_current_model = model
# Function tools for the manager agent to orchestrate the research process
@function_tool
async def plan_research(query: str) -> Dict[str, Any]:
""" Plan the research searches for a given query """
print("Planning searches...")
planner = create_planner_agent(_current_model)
result = await Runner.run(planner, f"Query: {query}")
search_plan = result.final_output_as(WebSearchPlan)
print(f"Will perform {len(search_plan.searches)} searches")
return {
"searches": [{"query": item.query, "reason": item.reason} for item in search_plan.searches],
"plan_ready": True
}
@function_tool
async def perform_search(search_query: str, reason: str) -> str:
""" Perform a single web search and return summarized results """
print(f"Searching for: {search_query}")
input_text = f"Search term: {search_query}\nReason for searching: {reason}"
try:
searcher = create_search_agent(_current_model)
result = await Runner.run(searcher, input_text)
return str(result.final_output)
except Exception as e:
print(f"Search failed for '{search_query}': {e}")
return f"Search failed for '{search_query}': {str(e)}"
@function_tool
async def write_initial_report(query: str, search_results: str) -> Dict[str, Any]:
""" Generate an initial research report from search results """
print("Writing initial report...")
input_text = f"Original query: {query}\nSummarized search results: {search_results}"
writer = create_writer_agent(_current_model)
result = await Runner.run(writer, input_text)
report_data = result.final_output_as(ReportData)
print("Initial report completed")
return {
"markdown_report": report_data.markdown_report,
"short_summary": report_data.short_summary,
"follow_up_questions": report_data.follow_up_questions
}
@function_tool
async def evaluate_report(query: str, report: str) -> Dict[str, Any]:
""" Evaluate the quality of a research report """
print("Evaluating report quality...")
input_text = f"Original Query: {query}\n\nReport to Evaluate:\n{report}"
evaluator = create_evaluator_agent(_current_model)
result = await Runner.run(evaluator, input_text)
evaluation = result.final_output_as(EvaluationResult)
print(f"Evaluation complete - Score: {evaluation.overall_score}/10, Needs refinement: {evaluation.needs_refinement}")
return {
"score": evaluation.overall_score,
"strengths": evaluation.strengths,
"weaknesses": evaluation.weaknesses,
"suggestions": evaluation.suggestions,
"needs_refinement": evaluation.needs_refinement,
"refinement_requirements": evaluation.refined_requirements
}
@function_tool
async def optimize_report(query: str, original_report: str, evaluation_feedback: str) -> str:
""" Optimize and improve a research report based on evaluation feedback """
print("Optimizing report...")
input_text = f"""Original Query: {query}
Original Report:
{original_report}
Evaluation Feedback:
{evaluation_feedback}
Please improve the report based on this feedback."""
optimizer = create_optimizer_agent(_current_model)
result = await Runner.run(optimizer, input_text)
optimized = result.final_output_as(OptimizedReport)
print("Report optimization complete")
return optimized.improved_markdown_report
# Regular function that can be called directly
async def _send_report_email_to_address(report: str, recipient_email: str) -> Dict[str, str]:
""" Send the final research report via email to a specific address """
import os
import sendgrid
from sendgrid.helpers.mail import Email, Mail, Content, To
print(f"Sending email to: {recipient_email}")
try:
sg = sendgrid.SendGridAPIClient(api_key=os.environ.get('SENDGRID_API_KEY'))
from_email = Email("mantomarchi300@outlook.com") # Verified sender
to_email = To(recipient_email) # User-provided email
# Create a nice subject line
subject = "🔍 Your Research Report is Ready"
# Convert markdown to HTML for better email formatting
import re
# Basic markdown to HTML conversion
html_report = report
# Convert markdown links to HTML links with styling
html_report = re.sub(r'\[([^\]]+)\]\(([^)]+)\)', r'\1', html_report)
# Convert headers
html_report = re.sub(r'^### (.*$)', r'
This report was created using advanced AI with multi-step quality assurance
Thank you for using our research service • Generated on {__import__('datetime').datetime.now().strftime('%B %d, %Y at %I:%M %p')}
"""
content = Content("text/html", html_content)
mail = Mail(from_email, to_email, subject, content).get()
response = sg.client.mail.send.post(request_body=mail)
print(f"Email response: {response.status_code}")
if response.status_code == 202:
return {"status": f"Email sent successfully to {recipient_email}"}
else:
return {"status": f"Email sending failed with status {response.status_code}"}
except Exception as e:
print(f"Email sending error: {e}")
return {"status": f"Email sending failed: {str(e)}"}
@function_tool
async def send_report_email_to_address(report: str, recipient_email: str) -> Dict[str, str]:
""" Send the final research report via email to a specific address """
return await _send_report_email_to_address(report, recipient_email)
@function_tool
async def send_report_email(report: str) -> Dict[str, str]:
""" Send the final research report via email (legacy function - uses hardcoded email) """
print("Sending email to default address...")
result = await Runner.run(email_agent, report)
print("Email sent to default address")
return {"status": "Email sent successfully to default address"}
# Manager Agent Instructions
MANAGER_INSTRUCTIONS = """
You are the Research Manager Agent responsible for orchestrating a comprehensive research process with quality assurance.
Your workflow:
1. **PLAN**: Use plan_research to create a search strategy for the query
2. **SEARCH**: Use perform_search for each planned search item to gather information
3. **INITIAL REPORT**: Use write_initial_report to create a first draft from all search results
4. **EVALUATE**: Use evaluate_report to assess the quality of the initial report
5. **OPTIMIZE** (if needed): If evaluation shows needs_refinement=True, use optimize_report to improve it
6. **FINALIZE**: Use send_report_email_to_address to deliver the final report
Quality Standards:
- Only proceed to email if the report scores 7+ or has been optimized
- If a report needs refinement, always optimize it before sending
- Ensure comprehensive coverage of the original query
- Maintain high standards for accuracy and completeness
Be methodical and ensure each step completes successfully before proceeding to the next.
"""
# Function to create custom research agent with email options
def create_custom_research_agent(email_address: str = None, send_email: bool = False, api_key: str = None, model: str = "gpt-4o-mini"):
"""Create a research manager agent with custom email settings"""
# Set API key in environment if provided
if api_key:
import os
os.environ["OPENAI_API_KEY"] = api_key
# Set the current model for all function tools
set_current_model(model)
if send_email and email_address:
# Include email sending in tools
tools = [
plan_research,
perform_search,
write_initial_report,
evaluate_report,
optimize_report,
send_report_email_to_address
]
instructions = f"""
You are the Research Manager Agent responsible for orchestrating a comprehensive research process with quality assurance.
Your workflow:
1. **PLAN**: Use plan_research to create a search strategy for the query
2. **SEARCH**: Use perform_search for each planned search item to gather information
3. **INITIAL REPORT**: Use write_initial_report to create a first draft from all search results
4. **EVALUATE**: Use evaluate_report to assess the quality of the initial report
5. **OPTIMIZE** (if needed): If evaluation shows needs_refinement=True, use optimize_report to improve it
6. **FINALIZE**: Use send_report_email_to_address with the report and recipient email "{email_address}" to deliver the final report
Quality Standards:
- Only proceed to email if the report scores 7+ or has been optimized
- If a report needs refinement, always optimize it before sending
- Ensure comprehensive coverage of the original query
- Maintain high standards for accuracy and completeness
IMPORTANT: When using send_report_email_to_address, you must provide both:
- The final report text as the first parameter
- The recipient email address "{email_address}" as the second parameter
Be methodical and ensure each step completes successfully before proceeding to the next.
The user has requested the report be emailed to: {email_address}
"""
else:
# Exclude email sending from tools
tools = [
plan_research,
perform_search,
write_initial_report,
evaluate_report,
optimize_report
]
instructions = """
You are the Research Manager Agent responsible for orchestrating a comprehensive research process with quality assurance.
Your workflow:
1. **PLAN**: Use plan_research to create a search strategy for the query
2. **SEARCH**: Use perform_search for each planned search item to gather information
3. **INITIAL REPORT**: Use write_initial_report to create a first draft from all search results
4. **EVALUATE**: Use evaluate_report to assess the quality of the initial report
5. **OPTIMIZE** (if needed): If evaluation shows needs_refinement=True, use optimize_report to improve it
6. **COMPLETE**: Return the final optimized report (do NOT send email - user chose not to receive email)
Quality Standards:
- Complete the process when report scores 7+ or has been optimized
- If a report needs refinement, always optimize it before completing
- Ensure comprehensive coverage of the original query
- Maintain high standards for accuracy and completeness
Be methodical and ensure each step completes successfully before proceeding to the next.
The user has chosen NOT to receive the report via email.
"""
return Agent(
name=f"Custom Research Manager Agent",
instructions=instructions,
tools=tools,
model=model,
handoff_description="Orchestrate comprehensive research with quality assurance and optional email delivery"
)
# Create the Research Manager Agent with agents-as-tools
ResearchManagerAgent = Agent(
name="Research Manager Agent",
instructions=MANAGER_INSTRUCTIONS,
tools=[
plan_research,
perform_search,
write_initial_report,
evaluate_report,
optimize_report,
send_report_email_to_address
],
model="gpt-4o-mini",
handoff_description="Orchestrate comprehensive research with quality assurance and optimization"
)
async def run_research_with_progress(query: str, email_address: str = None, send_email: bool = False, api_key: str = None, model: str = "gpt-4o-mini") -> AsyncGenerator[str, None]:
"""Run research with step-by-step progress updates"""
# Set API key if provided
if api_key:
import os
os.environ["OPENAI_API_KEY"] = api_key
# Set current model for function tools
set_current_model(model)
trace_id = gen_trace_id()
yield f"🚀 **Starting Enhanced Research**\n\n**Query:** {query}\n\n**Trace ID:** {trace_id}\n\n**Model:** {model}\n\n---\n\n"
try:
with trace("Enhanced Research with Progress", trace_id=trace_id):
# Step 1: Planning
yield "📋 **Step 1/6:** Planning research strategy...\n\n*Analyzing your query and determining the best search approach*"
planner = create_planner_agent(model)
result = await Runner.run(planner, f"Query: {query}")
search_plan = result.final_output_as(WebSearchPlan)
yield f"✅ **Planning Complete** - Will perform {len(search_plan.searches)} targeted searches\n\n---\n\n"
# Step 2: Searching
yield "🔍 **Step 2/6:** Conducting web searches...\n\n*Gathering information from multiple sources*"
search_results = []
for i, search_item in enumerate(search_plan.searches, 1):
yield f"🔍 **Search {i}/{len(search_plan.searches)}:** {search_item.query}\n\n*{search_item.reason}*"
try:
input_text = f"Search term: {search_item.query}\nReason for searching: {search_item.reason}"
searcher = create_search_agent(model)
result = await Runner.run(searcher, input_text)
search_results.append(str(result.final_output))
yield f"✅ **Search {i} Complete**\n\n"
except Exception as e:
yield f"⚠️ **Search {i} Failed:** {str(e)}\n\n"
search_results.append(f"Search failed: {str(e)}")
yield "✅ **All Searches Complete**\n\n---\n\n"
# Step 3: Writing Initial Report
yield "✍️ **Step 3/6:** Writing initial research report...\n\n*Analyzing and synthesizing all gathered information*"
combined_results = "\n\n".join(search_results)
input_text = f"Original query: {query}\nSummarized search results: {combined_results}"
writer = create_writer_agent(model)
result = await Runner.run(writer, input_text)
report_data = result.final_output_as(ReportData)
yield "✅ **Initial Report Complete**\n\n---\n\n"
# Step 4: Evaluating Quality
yield "🔍 **Step 4/6:** Evaluating report quality...\n\n*AI quality assessment in progress*"
input_text = f"Original Query: {query}\n\nReport to Evaluate:\n{report_data.markdown_report}"
evaluator = create_evaluator_agent(model)
result = await Runner.run(evaluator, input_text)
evaluation = result.final_output_as(EvaluationResult)
yield f"✅ **Evaluation Complete** - Score: {evaluation.overall_score}/10\n\n"
final_report = report_data.markdown_report
# Step 5: Optimization (if needed)
if evaluation.needs_refinement:
yield "🔧 **Step 5/6:** Optimizing report quality...\n\n*Improving report based on evaluation feedback*"
feedback = f"Score: {evaluation.overall_score}/10\nWeaknesses: {evaluation.weaknesses}\nSuggestions: {evaluation.suggestions}"
input_text = f"""Original Query: {query}
Original Report:
{report_data.markdown_report}
Evaluation Feedback:
{feedback}
Please improve the report based on this feedback."""
optimizer = create_optimizer_agent(model)
result = await Runner.run(optimizer, input_text)
optimized = result.final_output_as(OptimizedReport)
final_report = optimized.improved_markdown_report
yield "✅ **Optimization Complete** - Report quality improved\n\n---\n\n"
else:
yield "✅ **No Optimization Needed** - Report quality is excellent\n\n---\n\n"
# Step 6: Email Delivery (if requested)
if send_email and email_address:
yield f"📧 **Step 6/6:** Sending report to {email_address}...\n\n*Preparing and delivering your research report*"
try:
# Call the regular function directly
email_result = await _send_report_email_to_address(final_report, email_address)
email_status = email_result.get("status", "Email status unknown")
yield f"✅ **Email Delivery:** {email_status}\n\n---\n\n"
except Exception as e:
yield f"❌ **Email Failed:** {str(e)}\n\n---\n\n"
else:
yield "📄 **Step 6/6:** Finalizing report...\n\n*Email delivery not requested*\n\n---\n\n"
# Final result
yield f"""🎉 **Research Complete!**
**📊 Final Report:**
{final_report}
**🔗 View Full Trace:** https://platform.openai.com/traces/trace?trace_id={trace_id}
---
*Enhanced research completed with quality assurance*"""
except Exception as e:
yield f"❌ **Error during research:** {str(e)}"