from pydantic import BaseModel, Field
from agents import Agent
from gemini_model import gemini_model

HOW_MANY_SEARCHES = 10

INSTRUCTIONS = f"You are a helpful research assistant. Given a query (and potentially previous conversation history), \
come up with a set of web searches to perform to best answer the query.\n\n\
IMPORTANT CONTEXT AWARENESS:\n\
- If previous research is provided in the context, review it carefully\n\
- Plan searches that BUILD UPON previous findings rather than repeating them\n\
- Avoid redundant searches for topics already thoroughly covered\n\
- Focus on NEW aspects, follow-up questions, or gaps in previous research\n\
- If this is a follow-up question, tailor searches to specifically address the new angle\n\n\
Output {HOW_MANY_SEARCHES} diverse and comprehensive search terms that cover different aspects and perspectives, \
while avoiding redundancy with any previous research already conducted."


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.")

planner_agent = Agent(
    name="PlannerAgent",
    instructions=INSTRUCTIONS,
    model=gemini_model,
    output_type=WebSearchPlan,
)