agent.py

import os
import time
from google import genai
from google.genai import types

class GeminiAgent:
    def __init__(self, api_key=None, system_instruction=None):
        key = api_key or os.environ.get("GEMINI_API_KEY")
        if not key:
            raise ValueError("GEMINI_API_KEY is not set.")
        
        self.client = genai.Client(api_key=key)
        self.system_instruction = system_instruction
        self.models = [
            "gemini-2.5-flash",
            "gemini-2.5-pro",
            "gemini-2.0-flash",
            "gemini-2.0-flash-lite",
            "gemini-2.0-pro",
            "gemini-1.5-pro",
            "gemini-1.5-flash"
        ]

    def generate_content(self, prompt):
        config = types.GenerateContentConfig(
            system_instruction=self.system_instruction,
        )
        
        errors = {}
        for model in self.models:
            retries = 3
            for attempt in range(retries):
                try:
                    response = self.client.models.generate_content(
                        model=model,
                        contents=prompt,
                        config=config
                    )
                    return response.text
                except Exception as e:
                    error_msg = str(e)
                    # Nếu lỗi 429 (vượt quá quota), sẽ chờ một khoảng thời gian rồi gọi lại
                    if "429" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
                        wait_time = 15 * (attempt + 1)
                        print(f"[Warning] Model {model} hit Quota Limit (429). Retrying in {wait_time}s...")
                        time.sleep(wait_time)
                        errors[model] = f"429 Quota Exceeded (retried {attempt+1} times)"
                        continue
                    else:
                        # Các lỗi khác như 404 sẽ chuyển ngay sang model dự phòng
                        print(f"[Warning] Model {model} failed: {e}. Trying next model...")
                        errors[model] = error_msg
                        break
        
        raise RuntimeError(f"Tất cả model đều thất bại. Hãy kiểm tra lại Quota API của bạn. Chi tiết lỗi: {errors}")