luojunhui 1 месяц назад
Родитель
Сommit
f0b6aa6894
2 измененных файлов с 25 добавлено и 16 удалено
  1. 1 0
      applications/api/deep_seek_api_by_volcanoengine.py
  2. 24 16
      tasks/title_rewrite_task.py

+ 1 - 0
applications/api/deep_seek_api_by_volcanoengine.py

@@ -15,6 +15,7 @@ def fetch_deepseek_response(model, prompt):
     client = Ark(
         api_key=deep_seek_api_key_byte_dance,
         timeout=1800,
+        max_retries=2,
         )
     response = client.chat.completions.create(
         model=deep_seek_model.get(model, deep_seek_default_model),

+ 24 - 16
tasks/title_rewrite_task.py

@@ -1,6 +1,7 @@
 """
 @author: luojunhui
 """
+
 import traceback
 
 from pymysql.cursors import DictCursor
@@ -102,6 +103,7 @@ class TitleRewriteTask:
     """
     标题重写任务
     """
+
     def __init__(self):
         self.db = DatabaseConnector(db_config=long_articles_config)
         self.db.connect()
@@ -130,52 +132,61 @@ class TitleRewriteTask:
             set title_rewrite_status = %s
             where content_trace_id = %s and title_rewrite_status= %s;
         """
-        affected_rows = self.db.save(query=sql, params=(new_status, content_trace_id, ori_status))
+        affected_rows = self.db.save(
+            query=sql, params=(new_status, content_trace_id, ori_status)
+        )
         return affected_rows
 
     def insert_into_rewrite_table(self, content_trace_id, new_title):
         """
         insert into rewrite_table
         """
-        prompt_version = 'xx_250228'
+        prompt_version = "xx_250228"
         insert_sql = f"""
             insert into video_title_rewrite
             (content_trace_id, new_title, status, prompt_version)
             values (%s, %s, %s, %s);
         """
-        self.db.save(query=insert_sql, params=(content_trace_id, new_title, const.TITLE_USEFUL_STATUS, prompt_version))
+        self.db.save(
+            query=insert_sql,
+            params=(
+                content_trace_id,
+                new_title,
+                const.TITLE_USEFUL_STATUS,
+                prompt_version,
+            ),
+        )
 
     def rewrite_each_article(self, article):
         """
         rewrite each article
         """
-        content_trace_id = article['content_trace_id']
-        article_title = article['article_title']
+        content_trace_id = article["content_trace_id"]
+        article_title = article["article_title"]
 
         # lock each task
         affected_rows = self.update_title_rewrite_status(
             content_trace_id=content_trace_id,
             ori_status=const.TITLE_REWRITE_INIT_STATUS,
-            new_status=const.TITLE_REWRITE_LOCK_STATUS
+            new_status=const.TITLE_REWRITE_LOCK_STATUS,
         )
         if not affected_rows:
             return
 
         try:
             prompt = generate_prompt(article_title)
-            new_title = fetch_deepseek_response(model='default', prompt=prompt)
+            new_title = fetch_deepseek_response(model="default", prompt=prompt)
 
             # insert into rewrite table
             self.insert_into_rewrite_table(
-                content_trace_id=content_trace_id,
-                new_title=new_title
+                content_trace_id=content_trace_id, new_title=new_title
             )
 
             # unlock
             self.update_title_rewrite_status(
                 content_trace_id=content_trace_id,
                 ori_status=const.TITLE_REWRITE_LOCK_STATUS,
-                new_status=const.TITLE_REWRITE_SUCCESS_STATUS
+                new_status=const.TITLE_REWRITE_SUCCESS_STATUS,
             )
         except Exception as e:
             log(
@@ -186,13 +197,12 @@ class TitleRewriteTask:
                     "error_message": str(e),
                     "error_type": type(e).__name__,
                     "traceback": traceback.format_exc(),
-
-                }
+                },
             )
             self.update_title_rewrite_status(
                 content_trace_id=content_trace_id,
                 ori_status=const.TITLE_REWRITE_LOCK_STATUS,
-                new_status=const.TITLE_REWRITE_FAIL_STATUS
+                new_status=const.TITLE_REWRITE_FAIL_STATUS,
             )
 
     def deal(self):
@@ -203,6 +213,4 @@ class TitleRewriteTask:
         bar = tqdm(articles, desc="title rewrite task")
         for article in bar:
             self.rewrite_each_article(article)
-            bar.set_postfix({"content_trace_id": article['content_trace_id']})
-
-        
+            bar.set_postfix({"content_trace_id": article["content_trace_id"]})