from google import genai
from google.genai.types import Tool, GenerateContentConfig, GoogleSearch
import time
import os

api_key = os.getenv("GEMINI_API_KEY")
client = genai.Client(api_key=api_key)

model_id = "gemini-2.0-flash"

google_search_tool = Tool(
    google_search = GoogleSearch()
)

def webSearch(prompt):
  """
  Searches the web using Google Search.

  Args:
      prompt: A string representing the search query
      
  Returns:
      Search results in natural language.
  """  

  response = client.models.generate_content(
    model=model_id,
    contents=prompt,
    config=GenerateContentConfig(
        tools=[google_search_tool],
        response_modalities=["TEXT"],
    )
)
  
  # stream response
  for each in response.candidates[0].content.parts:
    for i in range(len(each.text)):
        time.sleep(0.001)
        yield each.text[: i+1]