File size: 854 Bytes
01f199c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | 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
|