OpenAI Research Agent

from agents import function_tool
import streamlit as st
from datetime import datetime

@function_tool
def save_important_fact(fact: str, source: str = None) -> str:
    """Save an important fact discovered during research.
    
    Args:
        fact: The important fact to save
        source: Optional source of the fact
    
    Returns:
        Confirmation message
    """
    if "collected_facts" not in st.session_state:
        st.session_state.collected_facts = []
    
    st.session_state.collected_facts.append({
        "fact": fact,
        "source": source or "Not specified",
        "timestamp": datetime.now().strftime("%H:%M:%S")
    })
    
    return f"Fact saved: {fact}"

from agents import Agent, Runner, WebSearchTool, handoff, trace
from pydantic import BaseModel

# Define data models
class ResearchPlan(BaseModel):
    topic: str
    search_queries: list[str]
    focus_areas: list[str]

class ResearchReport(BaseModel):
    title: str
    outline: list[str]
    report: str
    sources: list[str]
    word_count: int

# Research Agent
research_agent = Agent(
    name="Research Agent",
    instructions="""
    You are a research assistant. Given a search term, you search 
    the web and produce a concise summary of results. The summary 
    must be 2-3 paragraphs and less than 300 words.
    """,
    model="gpt-4o-mini",
    tools=[WebSearchTool(), save_important_fact],
)

# Editor Agent
editor_agent = Agent(
    name="Editor Agent",
    handoff_description="A senior researcher who writes comprehensive reports",
    instructions="""
    You are a senior researcher tasked with writing a cohesive report. 
    Create an outline first, then generate a lengthy and detailed report 
    in markdown format. Aim for 5-10 pages, at least 1000 words.
    """,
    model="gpt-4o-mini",
    output_type=ResearchReport,
)

# Triage Agent (Coordinator)
triage_agent = Agent(
    name="Triage Agent",
    instructions="""
    You are the coordinator. Your job is to:
    1. Understand the user's research topic
    2. Create a research plan (topic, search_queries, focus_areas)
    3. Hand off to the Research Agent
    4. After research, hand off to the Editor Agent
    """,
    handoffs=[handoff(research_agent), handoff(editor_agent)],
    model="gpt-4o-mini",
    output_type=ResearchPlan,
)

import asyncio
import uuid

async def run_research(topic):
    # Create a unique conversation ID for tracing
    conversation_id = str(uuid.uuid4().hex[:16])
    
    # Trace the entire workflow
    with trace("News Research", group_id=conversation_id):
        # Start with triage agent
        triage_result = await Runner.run(
            triage_agent,
            f"Research this topic thoroughly: {topic}"
        )
        
        # Get the research plan
        research_plan = triage_result.final_output
        
        # Editor agent compiles the report
        report_result = await Runner.run(
            editor_agent,
            triage_result.to_input_list()
        )
        
        return report_result.final_output

# Run the async function
report = asyncio.run(run_research("AI developments in 2024"))

import streamlit as st

# Initialize session state
if "conversation_id" not in st.session_state:
    st.session_state.conversation_id = str(uuid.uuid4().hex[:16])
if "collected_facts" not in st.session_state:
    st.session_state.collected_facts = []
if "research_done" not in st.session_state:
    st.session_state.research_done = False
if "report_result" not in st.session_state:
    st.session_state.report_result = None

from agents import trace

with trace("News Research", group_id=conversation_id):
    # All agent operations are traced
    result = await Runner.run(agent, message)

try:
    report_result = await Runner.run(editor_agent, triage_result.to_input_list())
    st.session_state.report_result = report_result.final_output
except Exception as e:
    st.error(f"Error generating report: {str(e)}")
    # Fallback to display raw agent response
    if hasattr(triage_result, 'new_items'):
        messages = [item for item in triage_result.new_items]
        raw_content = "\n\n".join([str(m.content) for m in messages])
        st.session_state.report_result = raw_content