from pydantic import BaseModel, Field
from agents import Agent

from agents import OpenAIChatCompletionsModel
from openai import AsyncOpenAI
import os

from dotenv import load_dotenv

load_dotenv(override=True)


model = "gemini-2.0-flash"
GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/"

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

client = AsyncOpenAI(base_url=GEMINI_BASE_URL,api_key=GEMINI_API_KEY)

custom_model = OpenAIChatCompletionsModel(model=model,openai_client=client)

HOW_MANY_SEARCHES = 5

INSTRUCTIONS = f"You are a helpful research assistant. Given a query, come up with a set of web searches \
to perform to best answer the query. Output {HOW_MANY_SEARCHES} terms to query for."


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=custom_model,
    output_type=WebSearchPlan,
)