planner_agent.py

from pydantic import BaseModel, Field
from agents import Agent
from typing import List

HOW_MANY_SEARCHES = 2


class QuestionAnswerPair(BaseModel):
    question: str = Field(description="The clarifying question that was asked")
    answer: str = Field(description="The user's answer to the clarifying question")


class ResearchContext(BaseModel):
    original_query: str = Field(description="The original research query from the user")
    clarification_qa: List[QuestionAnswerPair] = Field(
        description="List of question-answer pairs from the clarification process",
        min_items=3,
        max_items=3
    )


class WebSearchItem(BaseModel):
    reason: str = Field(description="Your reasoning for why this search is important to the query.")
    query: str = Field(description="The search term to use for the web search.")


class WebSearchPlan(BaseModel):
    searches: List[WebSearchItem] = Field(description="A list of web searches to perform to best answer the query.")
    
INSTRUCTIONS = f"""You are a helpful research assistant. Given a research context that includes:
1. The original user query
2. Three clarifying questions and their answers

Your task is to come up with {HOW_MANY_SEARCHES} web search terms that will best answer the user's query, 
taking into account both the original query AND the clarification answers provided.

The clarification answers provide important context about:
- What specific aspects the user is interested in
- The scope and depth they want
- Timeframe or other constraints
- Specific focus areas

Use this information to create more targeted and relevant search terms. Output exactly {HOW_MANY_SEARCHES} search terms."""

planner_agent = Agent(
    name="PlannerAgent",
    instructions=INSTRUCTIONS,
    model="gpt-4o-mini",
    output_type=WebSearchPlan,
)