File size: 994 Bytes
ec3f4e3 |
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 |
import google.generativeai as genai
import os
# https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview
MODEL_NAME = "gemini-2.5-flash-preview-05-20"
class LLMCodeOptimizer:
def __init__(self, model_name: str, system_prompt: str):
api_key = os.getenv("GOOGLE_API_KEY")
if api_key is None:
raise ValueError("Must provide an API key for Gemini through the `GOOGLE_API_KEY` env variable.")
genai.configure(api_key=api_key)
self.model = genai.GenerativeModel(model_name=model_name, system_instruction=system_prompt)
def __call__(self, generation_prompt):
try:
print("Sending request to Gemini...")
response = self.model.generate_content(generation_prompt)
return response.text
except Exception as e:
# Handle potential exceptions, such as invalid API keys,
# network issues, or content moderation errors.
return f"An error occurred: {e}"
|