src/models/Gemini.py

import pprint
import os
from google import genai
import dotenv
import time

from .Base import BaseModel

dotenv.load_dotenv()


class Gemini(BaseModel):
    def __init__(self, temperature=0):
        genai.configure(api_key=os.getenv("Google_API_KEY"))
        self.client = genai.GenerativeModel(model_name="gemini-2.5-flash")
        self.model = "gemini-2.5-flash"

    # @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(5))
    def prompt(self, processed_input):
        last_error = None
        for i in range(10):
            try:
                response = self.client.generate_content(
                    processed_input[0]['content']
                )
                return response.text, 0, 0
            except Exception as e:
                last_error = e
                time.sleep(2)

        return "", 0, 0