fix: handle google exceptions and add backoff

elleelleaime/evaluate/strategies/google/google.py

@@ -23,6 +23,9 @@ def _evaluate_impl(self, bug: Bug, sample: dict) -> Optional[List[dict]]:
 
         for generation in sample["generation"]:
             for candidate in generation["candidates"]:
+                if "content" not in candidate:
+                    evaluation.append(None)
+                    continue
                 candidate_patch = candidate["content"]["parts"][0]["text"]
                 candidate_patch = self.extract_patch_from_message(candidate_patch)
                 evaluation.append(

elleelleaime/generate/strategies/models/google/google.py

@@ -1,15 +1,23 @@
+import google.api_core
+import google.api_core.exceptions
 from elleelleaime.generate.strategies.strategy import PatchGenerationStrategy
 
 from dotenv import load_dotenv
 from typing import Any, List
 
 import os
+import tqdm
 import google.generativeai as genai
+import google
+import backoff
+
+import google.api
 
 
 class GoogleModels(PatchGenerationStrategy):
     def __init__(self, model_name: str, **kwargs) -> None:
         self.model_name = model_name
+        self.model = genai.GenerativeModel(self.model_name)
         self.temperature = kwargs.get("temperature", 0.0)
         self.n_samples = kwargs.get("n_samples", 1)
 
@@ -21,18 +29,20 @@ def __get_config(self):
             temperature=self.temperature,
         )
 
+    @backoff.on_exception(backoff.expo, google.api_core.exceptions.ResourceExhausted)
+    def __generate_with_backoff(self, prompt: str) -> dict:
+        completion = self.model.generate_content(
+            prompt, generation_config=self.__get_config()
+        )
+        return completion.to_dict()
+
     def _generate_impl(self, chunk: List[str]) -> Any:
         result = []
 
-        model = genai.GenerativeModel(self.model_name)
-
-        for prompt in chunk:
+        for prompt in tqdm.tqdm(chunk, "Generating patches for prompt..."):
             p_results = []
             for _ in range(self.n_samples):
-                completion = model.generate_content(
-                    prompt, generation_config=self.__get_config()
-                )
-                p_results.append(completion.to_dict())
+                p_results.append(self.__generate_with_backoff(prompt))
             result.append(p_results)
 
         return result