jan-contract / agents /general_assistant_agent.py
Amodit's picture
Force use of gemini-2.5-flash
188eb69
import os
import time
import random
from google import genai
from google.genai import types
# Configure the API key from the .env file
# Configure the API key from the .env file
try:
api_key = os.getenv("GOOGLE_API_KEY")
if api_key:
client = genai.Client(api_key=api_key)
model_name = "gemini-2.5-flash"
else:
print("CRITICAL WARNING: GOOGLE_API_KEY missing in general_assistant_agent.")
client = None
model_name = None
except Exception as e:
print(f"Error configuring Google Gen AI Client: {e}")
client = None
model_name = None
def ask_gemini(prompt: str) -> str:
"""
Sends a prompt directly to the Google Gen AI API using the new SDK.
Includes robust retry logic for 429 Resource Exhausted errors.
"""
if client is None:
return "Error: The Gen AI client is not configured. Please check your API key."
max_retries = 5
base_delay = 2 # seconds
for attempt in range(max_retries):
try:
response = client.models.generate_content(
model=model_name,
contents=prompt
)
return response.text
except Exception as e:
error_str = str(e)
if "429" in error_str or "RESOURCE_EXHAUSTED" in error_str:
if attempt == max_retries - 1:
return f"Error: Rate limit exceeded after {max_retries} attempts. Please try again later."
# Exponential backoff with jitter
delay = (base_delay * (2 ** attempt)) + random.uniform(0, 1)
print(f"Rate limit hit. Retrying in {delay:.2f} seconds...")
time.sleep(delay)
else:
return f"An error occurred while communicating with the Gemini API: {str(e)}"
return "Error: Failed to get response from Gemini API."