grounding_search.py

import os
import sys
import logging
import time
from typing import Optional
from dotenv import load_dotenv
from google import genai
from google.genai import types

load_dotenv('.env.local')  # Load environment variables from .env.local

def setup_logging() -> None:
    """Set up basic logging configuration."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s"
    )


def get_grounding_response(query: str, max_retries: int = 3) -> Optional[str]:
    """
    Generate content using Gemini with Google Search grounding.

    Args:
        query: The query string to process.
        max_retries: Maximum number of retries for transient errors.

    Returns:
        The generated response text, or None if an error occurs.
    """
    for attempt in range(max_retries + 1):
        try:
            # Initialize the client (assumes GOOGLE_API_KEY is set in environment)
            client = genai.Client()

            # Define the grounding tool
            grounding_tool = types.Tool(google_search=types.GoogleSearch())

            # Configure the generation settings
            config = types.GenerateContentConfig(tools=[grounding_tool])

            # Generate content
            response = client.models.generate_content(
                model="gemini-2.5-flash",
                contents=query,
                config=config,
            )

            return response.text
        except Exception as e:
            error_message = str(e)
            if "503" in error_message or "UNAVAILABLE" in error_message:
                if attempt < max_retries:
                    wait_time = 2 ** attempt  # Exponential backoff
                    logging.warning(f"Service unavailable, retrying in {wait_time} seconds... (attempt {attempt + 1}/{max_retries + 1})")
                    time.sleep(wait_time)
                    continue
            logging.error(f"An error occurred while generating content: {e}")
            return None
    return None


def main() -> None:
    """Main function to handle user input and generate response."""
    setup_logging()

    # Check for API key
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        logging.error("GEMINI_API_KEY environment variable is not set.")
        sys.exit(1)
    # Set GOOGLE_API_KEY for the library
    os.environ["GOOGLE_API_KEY"] = api_key

    # Get query from command line argument or prompt
    if len(sys.argv) > 1:
        query = " ".join(sys.argv[1:])
    else:
        query = input("Enter your query: ").strip()
        if not query:
            logging.error("No query provided.")
            sys.exit(1)

    logging.info(f"Processing query: {query}")

    # Get and print the response
    response = get_grounding_response(query)
    if response:
        print(response)
    else:
        logging.error("Failed to generate response.")


if __name__ == "__main__":
    main()