Alwrity - Bug fixes

lib/gpt_providers/text_generation/gemini_pro_text.py

@@ -4,6 +4,7 @@ import sys
 from pathlib import Path
 
 import google.generativeai as genai
+from google.api_core import retry
 from dotenv import load_dotenv
 load_dotenv(Path('../../../.env'))
 from loguru import logger
@@ -13,16 +14,10 @@ logger.add(sys.stdout,
         format="<level>{level}</level>|<green>{file}:{line}:{function}</green>| {message}"
     )
 
-from tenacity import (
-    retry,
-    stop_after_attempt,
-    wait_random_exponential,
-)  # for exponential backoff
 
-
-@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
 def gemini_text_response(prompt, temperature, top_p, n, max_tokens):
     """ Common functiont to get response from gemini pro Text. """
+    #FIXME: Include : https://github.com/google-gemini/cookbook/blob/main/quickstarts/rest/System_instructions_REST.ipynb
     try:
         genai.configure(api_key=os.getenv('GEMINI_API_KEY'))
     except Exception as err:
@@ -35,10 +30,13 @@ def gemini_text_response(prompt, temperature, top_p, n, max_tokens):
         "top_k": n,
         "max_output_tokens": max_tokens
     }
+    # FIXME: Expose model_name in main_config
     model = genai.GenerativeModel(model_name="gemini-1.0-pro", generation_config=generation_config)
     try:
-        response = model.generate_content(prompt, stream=True)
+        # text_response = []
+        response = model.generate_content(prompt, stream=True, request_options={'retry':retry.Retry()})
         for chunk in response:
+            # text_response.append(chunk.text)
             print(chunk.text)
         return response.text
     except Exception as err: