bugfix: fix minor bugs

requirements.txt

@@ -7,4 +7,4 @@ google-cloud-storage
 google-cloud-bigquery
 openai
 opencc
-gradio==4.13.0
+google-cloud-aiplatform

utils/completion_reward.py

@@ -21,11 +21,17 @@ from google.oauth2.service_account import Credentials
 from google import auth
 from google.cloud import bigquery
 from google.cloud import storage
+from google.cloud import aiplatform
+from vertexai.preview.generative_models import GenerativeModel
 
 SERVICE_ACCOUNT_INFO = os.getenv("GBQ_TOKEN")
 SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
 service_account_info_dict = json.loads(SERVICE_ACCOUNT_INFO)
 
+aiplatform.init(
+        project='junyiacademy',
+        service_account=service_account_info_dict,
+    )
 creds = Credentials.from_service_account_info(service_account_info_dict, scopes=SCOPES)
 
 gbq_client = bigquery.Client(
@@ -35,7 +41,6 @@ gcs_client = storage.Client(
     credentials=creds, project=service_account_info_dict["project_id"]
 )
 
-
 class CompletionReward:
     def __init__(self):
         self.player_backend_user_id = None
@@ -262,6 +267,7 @@ class OpenAIAgent:
         retry_attempts = 0
         while retry_attempts < 5:
             try:
+                logging.info("Generating image...")
                 response = client.images.generate(
                     model="dall-e-3",
                     prompt="Create an image in a retro Ghibli style, with a focus on a universe theme. The artwork should maintain the traditional hand-drawn animation look characteristic of Ghibli. Imagine a scene set in outer space or a fantastical cosmic environment, rich with vibrant and varied color palettes to capture the mystery and majesty of the universe. The background should be detailed, showcasing stars, planets, and nebulae, blending the Ghibli style's nostalgia and emotional depth with the awe-inspiring aspects of space. The overall feel should be timeless, merging the natural wonder of the cosmos with the storytelling and emotional resonance typical of the retro Ghibli aesthetic. Soft lighting and gentle shading should be used to enhance the dreamlike, otherworldly quality of the scene.",
@@ -275,7 +281,7 @@ class OpenAIAgent:
 
             except Exception as e:
                 retry_attempts += 1
-                logging.error(f"Attempt {retry_attempts}: {e}")
+                logging.error(f"DALLE Attempt {retry_attempts}: {e}")
                 time.sleep(1 * retry_attempts)  # exponential backoff
 
 
@@ -319,15 +325,6 @@ class AWSAgent:
 
 class GoogleAgent:
     def get_story(self, user_log):
-        creds.refresh(Request())
-
-        url = f"https://us-central1-aiplatform.googleapis.com/v1/projects/junyiacademy/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent?"
-
-        # Headers
-        headers = {
-            "Authorization": f"Bearer {creds.token}",
-            "Content-Type": "application/json",
-        }
         system_prompt = """
                     我正在舉辦一個學習型的活動，參與活動的學生為 1-9 年級，我為他們設計了一個獨特的故事機制，每天每個學生都會收到屬於自己獨特的冒險紀錄，現在我需要你協助我將這些冒險紀錄，製作成一段冒險故事，請
                     - 以「你」稱呼學生
@@ -342,36 +339,17 @@ class GoogleAgent:
             ```
         """
 
-        # Provided payload
-        payload = {
-            "contents": [
-                {
-                    "role": "USER",
-                    "parts": {"text": f"{system_prompt}, 以下是我的冒險故事 ```{user_log}```"},
-                },
-            ],
-            "safety_settings": {
-                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-                "threshold": "BLOCK_LOW_AND_ABOVE",
-            },
-            "generation_config": {
-                "temperature": 0.2,
-                "topP": 0.8,
-                "topK": 40,
-                "maxOutputTokens": 2048,
-            },
-        }
+        gemini_pro_model = GenerativeModel("gemini-pro")
 
         retry_attempts = 0
         while retry_attempts < 5:
             try:
-                response = requests.post(url, headers=headers, data=json.dumps(payload))
+                logging.info("Google Generating response...")
+                model_response = gemini_pro_model.generate_content(f"{system_prompt}, 以下是我的冒險故事 ```{user_log}```")
+                
                 chinese_converter = OpenCC("s2tw")
 
-                final_story = ""
-                for i in response.json():
-                    final_story += i["candidates"][0]["content"]["parts"][0]["text"]
-                return chinese_converter.convert(final_story)
+                return chinese_converter.convert(model_response.candidates[0].content.parts[0].text)
 
             except Exception as e:
                 retry_attempts += 1