dev

applications/api/google_ai_api.py

@@ -11,8 +11,8 @@ class GoogleAIAPI(object):
     """
     Google 视频内容理解API
     """
-    def __init__(self):
-        self.api_key = 'AIzaSyBAFFL6yHa4kUK-YcuAw8tbiSHQB6oJG34'
+    def __init__(self, api_key):
+        self.api_key = api_key
         self.client = genai.Client(api_key=self.api_key)
 
     def upload_file(self, file_path: str):
@@ -64,16 +64,27 @@ class GoogleAIAPI(object):
         """
         try:
             response = self.client.models.generate_content(
-                model='gemini-2.0-flash',
+                model='gemini-1.5-flash',
                 contents=[
                     video_file,
                     prompt
                 ]
             )
-            return response.text
+            return {
+                'code': 200,
+                'text': response.text
+            }
         except ClientError as e:
-            print(e.response.status_code)
-            return None
+            return {
+                'code': e.response.status_code,
+                'text': 'resource exhausted'
+            }
+        except Exception as e:
+            return {
+                'code': 500,
+                'text': str(e)
+            }
+
 
     def delete_video(self, file_name: str):
         """

coldStartTasks/ai_pipeline/extract_video_best_frame.py

@@ -22,8 +22,7 @@ from coldStartTasks.ai_pipeline.basic import get_video_cover
 from coldStartTasks.ai_pipeline.basic import normalize_time_str
 
 const = GoogleVideoUnderstandTaskConst()
-google_ai = GoogleAIAPI()
-
+google_ai = GoogleAIAPI(api_key='AIzaSyBAFFL6yHa4kUK-YcuAw8tbiSHQB6oJG34')
 
 class ExtractVideoBestFrame:
     """
@@ -33,6 +32,7 @@ class ExtractVideoBestFrame:
     def __init__(self):
         self.db_client = DatabaseConnector(db_config=long_articles_config)
         self.db_client.connect()
+        self.api_status = True
 
     def _roll_back_lock_tasks(self, task: str) -> int:
         return roll_back_lock_tasks(
@@ -222,6 +222,14 @@ class ExtractVideoBestFrame:
     def upload_video_to_gemini_ai(
         self, max_processing_pool_size: int = const.POOL_SIZE
     ) -> None:
+        if not self.api_status:
+            log(
+                task=const.TASK_NAME,
+                function="upload_video_to_gemini_ai",
+                message="api_status is False, skip this task",
+            )
+            return
+
         # upload video to gemini ai
         roll_back_lock_tasks_count = self._roll_back_lock_tasks(task="upload")
         log(
@@ -305,19 +313,98 @@ class ExtractVideoBestFrame:
 
                 case "ACTIVE":
                     # video process successfully
-                    # try:
-                    best_frame_time_ms = google_ai.fetch_info_from_google_ai(
-                        prompt=extract_best_frame_prompt(),
-                        video_file=google_file,
-                    )
-                    print(best_frame_time_ms)
-                    if best_frame_time_ms:
-                        self.set_extract_result(
-                            task_id=task["id"],
-                            file_state="ACTIVE",
-                            best_frame_time_ms=best_frame_time_ms.strip(),
+                    try:
+                        response = google_ai.fetch_info_from_google_ai(
+                            prompt=extract_best_frame_prompt(),
+                            video_file=google_file,
+                        )
+                        code = response['code']
+                        match code:
+                            case 200:
+                                best_frame_time_ms = response['text']
+                                if best_frame_time_ms:
+                                    self.set_extract_result(
+                                        task_id=task["id"],
+                                        file_state="ACTIVE",
+                                        best_frame_time_ms=best_frame_time_ms.strip(),
+                                    )
+                                else:
+                                    update_task_queue_status(
+                                        db_client=self.db_client,
+                                        task_id=task["id"],
+                                        task="extract",
+                                        ori_status=const.PROCESSING_STATUS,
+                                        new_status=const.FAIL_STATUS,
+                                    )
+                                # delete local file and google file
+                                if os.path.exists(video_local_path):
+                                    os.remove(video_local_path)
+
+                                google_ai.delete_video(file_name)
+                                log(
+                                    task=const.TASK_NAME,
+                                    function="extract_best_frame_with_gemini_ai",
+                                    message="video process successfully",
+                                    data={
+                                        "task_id": task["id"],
+                                        "file_name": file_name,
+                                        "state": state,
+                                        "best_frame_time_ms": best_frame_time_ms,
+                                    },
+                                )
+                            case 429:
+                                # resource exhausted
+                                log(
+                                    task=const.TASK_NAME,
+                                    function="extract_best_frame_with_gemini_ai",
+                                    message="resource exhausted",
+                                    data={
+                                        "task_id": task["id"],
+                                        "file_name": file_name,
+                                        "state": state,
+                                        "response": response
+                                    },
+                                )
+                                update_task_queue_status(
+                                    db_client=self.db_client,
+                                    task_id=task["id"],
+                                    task="extract",
+                                    ori_status=const.PROCESSING_STATUS,
+                                    new_status=const.INIT_STATUS
+                                )
+                                return
+                            case 500:
+                                # other error
+                                log(
+                                    task=const.TASK_NAME,
+                                    function="extract_best_frame_with_gemini_ai",
+                                    message="other error",
+                                    data={
+                                        "task_id": task["id"],
+                                        "file_name": file_name,
+                                        "state": state,
+                                        "response": response,
+                                    },
+                                )
+                                update_task_queue_status(
+                                    db_client=self.db_client,
+                                    task_id=task["id"],
+                                    task="extract",
+                                    ori_status=const.PROCESSING_STATUS,
+                                    new_status=const.FAIL_STATUS
+                                )
+
+                    except Exception as e:
+                        log(
+                            task=const.TASK_NAME,
+                            function="extract_best_frame_with_gemini_ai",
+                            message="task_failed_inside_cycle",
+                            data={
+                                "task_id": task["id"],
+                                "track_back": traceback.format_exc(),
+                                "error": str(e),
+                            },
                         )
-                    else:
                         update_task_queue_status(
                             db_client=self.db_client,
                             task_id=task["id"],
@@ -325,40 +412,6 @@ class ExtractVideoBestFrame:
                             ori_status=const.PROCESSING_STATUS,
                             new_status=const.FAIL_STATUS,
                         )
-                    # delete local file and google file
-                    if os.path.exists(video_local_path):
-                        os.remove(video_local_path)
-
-                    google_ai.delete_video(file_name)
-                    log(
-                        task=const.TASK_NAME,
-                        function="extract_best_frame_with_gemini_ai",
-                        message="video process successfully",
-                        data={
-                            "task_id": task["id"],
-                            "file_name": file_name,
-                            "state": state,
-                            "best_frame_time_ms": best_frame_time_ms,
-                        },
-                    )
-                    # except Exception as e:
-                    #     log(
-                    #         task=const.TASK_NAME,
-                    #         function="extract_best_frame_with_gemini_ai",
-                    #         message="task_failed_inside_cycle",
-                    #         data={
-                    #             "task_id": task["id"],
-                    #             "track_back": traceback.format_exc(),
-                    #             "error": str(e),
-                    #         },
-                    #     )
-                    #     update_task_queue_status(
-                    #         db_client=self.db_client,
-                    #         task_id=task["id"],
-                    #         task="extract",
-                    #         ori_status=const.PROCESSING_STATUS,
-                    #         new_status=const.FAIL_STATUS,
-                    #     )
 
         except Exception as e:
             log(
@@ -380,6 +433,14 @@ class ExtractVideoBestFrame:
             )
 
     def extract_best_frame_with_gemini_ai(self):
+        if not self.api_status:
+            log(
+                task=const.TASK_NAME,
+                function="extract_best_frame_with_gemini_ai",
+                message="api_status is False, skip this task",
+            )
+            return
+
         # roll back lock tasks
         roll_back_lock_tasks_count = self._roll_back_lock_tasks(task="extract")
         log(
@@ -447,6 +508,13 @@ class ExtractVideoBestFrame:
         """
         get cover with best frame
         """
+        if not self.api_status:
+            log(
+                task=const.TASK_NAME,
+                function="extract_cover_with_ffmpeg",
+                message="api_status is False, skip this task",
+            )
+            return
         # roll back lock tasks
         roll_back_lock_tasks_count = self._roll_back_lock_tasks(task="get_cover")
         log(