Merge branch '2025-04-23-add-category-to-cold-start-pool' of luojunhui/LongArticlesJob into master

applications/api/deep_seek_api_official.py

@@ -0,0 +1,40 @@
+"""
+deep_seek 官方 API
+"""
+import json
+from openai import OpenAI
+
+from config import deep_seek_official_model
+from config import deep_seek_official_api_key
+
+def fetch_deepseek_completion(model, prompt, output_type='text'):
+    """
+    deep_seek方法
+    """
+    client = OpenAI(
+        api_key=deep_seek_official_api_key,
+        base_url="https://api.deepseek.com"
+    )
+
+    # get response format
+    if output_type == "json":
+        response_format = {"type": "json_object"}
+    else:
+        response_format = {"type": "text"}
+
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": prompt,
+            }
+        ],
+        model=deep_seek_official_model.get(model, "deepseek-chat"),
+        response_format=response_format,
+    )
+    response = chat_completion.choices[0].message.content
+    if output_type == "json":
+        response_json = json.loads(response)
+        return response_json
+
+    return response

applications/const/__init__.py

@@ -455,3 +455,33 @@ class GoogleVideoUnderstandTaskConst:
     TASK_NAME = "extract_video_best_frame_as_cover"
     DIR_NAME = "static"
     POOL_SIZE = 15
+
+
+class CategoryGenerationTaskConst:
+    """
+    const for category generation task
+    """
+    # MAX THREAD
+    MAX_WORKERS = 5
+
+    # task batch size
+    BATCH_SIZE = 20
+
+    # min batch
+    MIN_BATCH_SIZE = 1
+
+    # article status
+    ARTICLE_GOOD_STATUS = 0
+
+    # task status
+    INIT_STATUS = 0
+    PROCESSING_STATUS = 1
+    SUCCESS_STATUS = 2
+    FAIL_STATUS = 99
+
+    # max processing time
+    MAX_PROCESSING_TIME = 3600
+
+    # task info
+    TABLE_NAME = "publish_single_video_source"
+    TASK_NAME = "generate_category_with_title"

applications/utils/common.py

@@ -58,4 +58,14 @@ def request_retry(retry_times, min_retry_delay, max_retry_delay):
         retry=retry_if_exception_type((RequestException, TimeoutError)),
         reraise=True  # 重试耗尽后重新抛出异常
     )
-    return common_retry
+    return common_retry
+
+def yield_batch(data, batch_size):
+    """
+    生成批次数据
+    :param data:
+    :param batch_size:
+    :return:
+    """
+    for i in range(0, len(data), batch_size):
+        yield data[i:i + batch_size]

config/__init__.py

@@ -90,6 +90,15 @@ moon_shot = {
 }
 
 
+# deepseek official api
+deep_seek_official_api_key = 'sk-cfd2df92c8864ab999d66a615ee812c5'
+
+deep_seek_official_model = {
+    "DeepSeek-R1": "deepseek-reasoner",
+    "DeepSeek-V3": "deepseek-chat",
+}
+
+# deepseek volcano engine
 deep_seek_model = {
     "DeepSeek-R1": "ep-20250213194143-d8q4t",
     "DeepSeek-V3": "ep-20250213194558-rrmr2"

run_title_rewrite_task.py

@@ -1,9 +0,0 @@
-"""
-@author: luojunhui
-"""
-from tasks.title_rewrite_task import TitleRewriteTask
-
-
-if __name__ == '__main__':
-    task = TitleRewriteTask()
-    task.deal()

sh/run_title_rewrite_task.sh → sh/run_title_process_task.sh

@@ -4,15 +4,15 @@
 CURRENT_DATE=$(date +%F)
 
 # 日志文件路径，包含日期
-LOG_FILE="/root/luojunhui/logs/title_rewrite_task_log_$CURRENT_DATE.txt"
+LOG_FILE="/root/luojunhui/logs/title_process_task_log_$CURRENT_DATE.txt"
 
 # 重定向整个脚本的输出到带日期的日志文件
 exec >> "$LOG_FILE" 2>&1
-if pgrep -f "python3 run_title_rewrite_task.py" > /dev/null
+if pgrep -f "python3 title_process_task.py" > /dev/null
 then
-    echo "$(date '+%Y-%m-%d %H:%M:%S') - run_title_rewrite_task.py is running"
+    echo "$(date '+%Y-%m-%d %H:%M:%S') - title_process_task.py is running"
 else
-    echo "$(date '+%Y-%m-%d %H:%M:%S') - trying to restart run_title_rewrite_task.py"
+    echo "$(date '+%Y-%m-%d %H:%M:%S') - trying to restart title_process_task.py"
     # 切换到指定目录
     cd /root/luojunhui/LongArticlesJob
 
@@ -21,6 +21,6 @@ else
     conda activate tasks
 
     # 在后台运行 Python 脚本并重定向日志输出
-    nohup python3 run_title_rewrite_task.py >> "${LOG_FILE}" 2>&1 &
-    echo "$(date '+%Y-%m-%d %H:%M:%S') - successfully restarted run_title_rewrite_task.py"
+    nohup python3 title_process_task.py >> "${LOG_FILE}" 2>&1 &
+    echo "$(date '+%Y-%m-%d %H:%M:%S') - successfully restarted title_process_task.py"
 fi

tasks/ai_tasks/category_generation_task.py

@@ -0,0 +1,241 @@
+"""
+generate category for given title
+"""
+
+import time
+import concurrent
+import traceback
+from concurrent.futures import ThreadPoolExecutor
+
+from pymysql.cursors import DictCursor
+from tqdm import tqdm
+
+from applications import log
+from applications.api.deep_seek_api_official import fetch_deepseek_completion
+from applications.const import CategoryGenerationTaskConst
+from applications.db import DatabaseConnector
+from applications.utils import yield_batch
+
+from config import long_articles_config
+from tasks.ai_tasks.prompts import category_generation_from_title
+
+
+class CategoryGenerationTask:
+
+    def __init__(self):
+        self.db_client = DatabaseConnector(long_articles_config)
+        self.db_client.connect()
+        self.const = CategoryGenerationTaskConst()
+
+    def set_category_status_as_success(
+        self, thread_db_client: DatabaseConnector, article_id: int, category: str
+    ) -> int:
+
+        update_query = f"""
+            update publish_single_video_source
+            set category = %s, category_status = %s, category_status_update_ts = %s
+            where id = %s and category_status = %s;
+        """
+
+        update_rows = thread_db_client.save(
+            query=update_query,
+            params=(
+                category,
+                self.const.SUCCESS_STATUS,
+                int(time.time()),
+                article_id,
+                self.const.PROCESSING_STATUS,
+            ),
+        )
+        return update_rows
+
+    def set_category_status_as_fail(
+        self, thread_db_client: DatabaseConnector, article_id: int
+    ) -> int:
+
+        update_query = f"""
+            update publish_single_video_source
+            set category_status = %s, category_status_update_ts = %s
+            where id = %s and category_status = %s;
+        """
+
+        update_rows = thread_db_client.save(
+            query=update_query,
+            params=(
+                self.const.FAIL_STATUS,
+                int(time.time()),
+                article_id,
+                self.const.PROCESSING_STATUS,
+            ),
+        )
+        return update_rows
+
+    def update_title_category(
+        self, thread_db_client: DatabaseConnector, article_id: int, completion: dict
+    ):
+
+        try:
+            category = completion.get(str(article_id))
+            self.set_category_status_as_success(thread_db_client, article_id, category)
+
+        except Exception as e:
+            log(
+                task=self.const.TASK_NAME,
+                function="update_each_record_status",
+                message="AI返回格式失败，更新状态为失败",
+                data={
+                    "article_id": article_id,
+                    "error": str(e),
+                    "traceback": traceback.format_exc(),
+                },
+            )
+            self.set_category_status_as_fail(thread_db_client, article_id)
+
+    def rollback_lock_tasks(self) -> int:
+
+        update_query = f"""
+            update publish_single_video_source  
+            set category_status = %s
+            where category_status = %s and category_status_update_ts <= %s;
+        """
+
+        update_rows = self.db_client.save(
+            query=update_query,
+            params=(
+                self.const.INIT_STATUS,
+                self.const.PROCESSING_STATUS,
+                int(time.time()) - self.const.MAX_PROCESSING_TIME,
+            ),
+        )
+
+        return update_rows
+
+    def lock_task(
+        self, thread_db_client: DatabaseConnector, article_id_tuple: tuple[int, ...]
+    ) -> int:
+
+        update_query = f"""
+            update publish_single_video_source
+            set category_status = %s, category_status_update_ts = %s
+            where id in %s and category_status = %s;
+        """
+
+        update_rows = thread_db_client.save(
+            query=update_query,
+            params=(
+                self.const.PROCESSING_STATUS,
+                int(time.time()),
+                article_id_tuple,
+                self.const.INIT_STATUS,
+            ),
+        )
+        return update_rows
+
+    def deal_each_article(self, thread_db_client, article: dict):
+        """
+        deal each article
+        """
+        article_id = article["id"]
+        title = article["article_title"]
+
+        title_batch = [(article_id, title)]
+
+        prompt = category_generation_from_title(title_batch)
+        try:
+            completion = fetch_deepseek_completion(
+                model="DeepSeek-V3", prompt=prompt, output_type="json"
+            )
+            self.update_title_category(thread_db_client, article_id, completion)
+
+        except Exception as e:
+            log(
+                task=self.const.TASK_NAME,
+                message="该文章存在敏感词，AI 拒绝返回",
+                function="deal_each_article",
+                data={
+                    "article_id": article_id,
+                    "error": str(e),
+                    "traceback": traceback.format_exc(),
+                },
+            )
+            self.set_category_status_as_fail(thread_db_client, article_id)
+
+    def deal_batch_in_each_thread(self, task_batch: list[dict]):
+        """
+        deal in each thread
+        """
+        thread_db_client = DatabaseConnector(long_articles_config)
+        thread_db_client.connect()
+
+        title_batch = [(i["id"], i["article_title"]) for i in task_batch]
+        id_tuple = tuple([int(i["id"]) for i in task_batch])
+
+        lock_rows = self.lock_task(thread_db_client, id_tuple)
+        if lock_rows:
+            prompt = category_generation_from_title(title_batch)
+
+            try:
+                completion = fetch_deepseek_completion(
+                    model="DeepSeek-V3", prompt=prompt, output_type="json"
+                )
+            except Exception as e:
+                log(
+                    task=self.const.TASK_NAME,
+                    function="category_generation_task",
+                    message=" batch 中存在敏感词，AI 拒绝返回",
+                    data={
+                        "article_id": id_tuple,
+                        "error": str(e),
+                        "traceback": traceback.format_exc(),
+                    },
+                )
+                for article in tqdm(task_batch):
+                    self.deal_each_article(thread_db_client, article)
+
+                return
+
+            for article in title_batch:
+                self.update_title_category(thread_db_client, article[0], completion)
+
+        else:
+            return
+
+    def get_task_list(self):
+        """
+        get task_list from a database
+        """
+        fetch_query = f"""
+            select id, article_title from publish_single_video_source
+            where category_status = %s and bad_status = %s
+            order by score desc;
+        """
+        fetch_result = self.db_client.fetch(
+            query=fetch_query,
+            cursor_type=DictCursor,
+            params=(self.const.INIT_STATUS, self.const.ARTICLE_GOOD_STATUS),
+        )
+        return fetch_result
+
+    def deal(self):
+
+        self.rollback_lock_tasks()
+
+        task_list = self.get_task_list()
+        task_batch_list = yield_batch(data=task_list, batch_size=self.const.BATCH_SIZE)
+
+        ##  dev
+        # for task_batch in task_batch_list:
+        #     self.deal_batch_in_each_thread(task_batch)
+
+        with ThreadPoolExecutor(max_workers=self.const.MAX_WORKERS) as executor:
+            futures = [
+                executor.submit(self.deal_batch_in_each_thread, task_batch)
+                for task_batch in task_batch_list
+            ]
+
+            for _ in tqdm(
+                concurrent.futures.as_completed(futures),
+                total=len(futures),
+                desc="Processing batches",
+            ):
+                pass

tasks/ai_tasks/prompts.py

@@ -0,0 +1,127 @@
+"""
+ai tasks prompt
+"""
+
+def category_generation_from_title(title_list):
+    """
+    generate prompt category for given title
+    """
+    prompt = f"""
+        请帮我完成以下任务：输入为文章的标题，根据标题判断其内容所属的类目，输出为文章标题及其对应的类目。
+        类目需从以下15个品类内选择：
+        1. 知识科普
+        定义：以通俗易懂的方式普及科学、技术、健康、安全、生活常识、财产保护、医保政策、为人处事方式等内容，旨在提高公众的知识水平和认知能力。内容通常具有教育性和实用性，涵盖自然、社会、文化等多个领域。
+        标题示例：
+        我国存款最安全的五大银行，永远都不会倒闭，你知道是哪五家吗？
+        借条上不要写“这3个字”，不然变成一张废纸，否则用法律也没用
+        不能二次加热的3种食物！再次提醒：这3种食物吃不完最好扔掉
+        
+        2. 军事历史
+        定义：聚焦于历史上的军事事件、战争故事、军事策略、英雄人物等内容，旨在还原战争场景、探讨军事决策、揭示历史真相，并展现战争中的人物命运与历史影响。内容通常以叙事、分析或回忆的形式呈现，兼具历史深度和故事性。
+        标题示例：
+        对越作战永远失踪的332人，陵园没有墓碑，没有名字，只有烈士证
+        淮海大战丢失阵地，师长带头冲锋！最后出一口恶气：活捉敌最高指挥官
+        抗战时，一村民被敌拉去带路，半道回头忽发现:后面跟个游击队员
+        
+        3. 家长里短
+        定义：围绕家庭成员之间的关系、矛盾、情感、道德、等展开的故事或讨论，内容常涉及婚姻、亲子、婆媳、兄弟姐妹等关系，或是人情往来、金钱纠纷、情感变化等内容，反映家庭生活中的温情、冲突与人性。
+        标题示例：
+        父母越老越能暴露家庭最真实的一面：当父母70岁，子女不该抱有这三种期待
+        老母亲分家产，给亲闺女30万，给养女一筐青菜，养女意外摔倒，看到筐子里的东西，瞬间愣住了
+        我花150一天雇了阿姨，两天后上班回来给她300，阿姨说我账算错了
+        
+        4. 社会法治
+        定义：聚焦社会事件、法律纠纷、法院判决、社会现象等内容，通常涉及道德、法律、公平正义等议题，旨在揭示社会问题、探讨法律规则或反映人性与社会现实。
+        标题示例：
+        山东，女子在小区捡到16万天价项链，业主悬赏3万找回，女子归还后，失主拒绝支付报酬，还说：要有格局，女子认为被骗，将失主告上法庭
+        陕西，女子22万买26层房，2年后，楼盘24层就已经封顶！开发商：你闹事造成100万损失，道歉才给赔偿！
+        上海：男子超市连续购买46枚过期咸鸭蛋，2天分46次交易，向厂家索赔金14万，法院判了！
+        
+        5. 奇闻趣事
+        定义：以猎奇、娱乐为主，涵盖罕见、奇特、有趣的事件、发现或故事，内容通常具有趣味性和话题性，能够引发读者的好奇心和讨论。
+        标题示例：
+        狗屎运？江西男子钓鱼时发现青鱼尸骸，扒开后捡到鸡蛋大小的青鱼石，网友：起码值几千！
+        内蒙古小伙河边捡到金牌，拒绝上交将其熔成金手镯，专家气愤不已
+        男子买了一辆废弃坦克，拆油箱时，他发现了一根又一根的金条……
+        
+        6. 名人八卦
+        定义：围绕名人的生活、言论、事件、八卦等内容展开，通常涉及娱乐圈、政界、历史人物等，旨在满足公众对名人隐私和动态的好奇心。
+        标题示例：
+        难怪王扶林说陈晓旭不够漂亮，看看他选的原黛玉候选人，那才叫美
+        心狠手辣的容嬷嬷年轻时是校花？看了照片后，网友直接闭嘴了！
+        李玉成终于说出实话，公开吐槽马玉琴年纪太大，结婚28年疑似后悔
+        
+        7. 健康养生
+        定义：关注健康、养生、疾病预防、生活习惯等方面的知识和建议，内容通常具有实用性和指导性，旨在帮助读者改善生活质量、提升健康水平。
+        标题示例：
+        72岁老人每天一个蒸苹果，半年后体检，看到指标变化让他乐开了花
+        40岁女子每天吃水煮蛋，一年后去体检，检查报告令医生都羡慕不已
+        2024年血糖新标准已公布，不再是3.9~6.1，你的血糖还不算高吗？
+        
+        8. 情感故事
+        定义：以人与人之间的情感交流、感人故事、情感经历为主题，内容通常充满温情、感动或反思，旨在引发读者的情感共鸣和思考。
+        标题示例：
+        男孩饭店吃饭，发现陌生女子和去世母亲很像，走过去说：我妈妈去世了，能抱一下我吗？
+        河南一女子直播时，被失散 32 年的父亲认出：闺女等着爸爸接你回家
+        1987年，江苏男子借好友一千元，25年后朋友成富豪还他1000万报恩
+        流浪狗跟着骑行夫妻跑了一百多公里，一直守护在女主身边，赶都赶不走，当男主得知原因后竟抱着狗狗大哭起来
+        
+        9. 国家大事
+        定义：涉及国家实力、科技发展、资源发现、国际合作等内容，通常以宏观视角展现国家的综合实力、科技成就或国际影响力，体现国家的崛起与发展。
+        标题示例：
+        我国在南极发现“海上粮仓”，储量高达10亿吨，世界各国眼红不已
+        我国贵州发现7000万吨宝藏，价值高达上万亿，多国求合作被拒绝
+        距我国3000公里，塞班岛明明归美国管辖，为何岛上大多是中国人？
+        
+        10. 现代人物
+        定义：聚焦活跃在21世纪后具有传奇色彩或巨大贡献的人物、事迹、成就等，内容通常充满戏剧性和启发性，旨在展现人物的非凡经历或历史贡献。
+        标题示例：
+        她曾狂贪国家上百亿，被发现时已经移居美国，最终还风光一时得善终
+        山东女子因坐月子无聊，破译美国2套绝密系统的密码，国家：奖励711万！
+        牺牲太大了！航天女英雄刘洋：结婚8年未生子，回地面后“消失”的她怎样了？
+        
+        11. 怀旧时光
+        定义：以回忆和怀旧为主题，涉及过去的历史、文化、生活、照片等内容，旨在唤起读者对过去时光的情感共鸣和怀念。
+        标题示例：
+        1975年“下馆子”的老照片，2元能吃些什么，勾起那段最难忘的时光
+        82年，北京老人捡回两张“破椅子”，遭家人数落，29年后拍出2300万
+        这张老照片第一次看到，邓颖超和李讷的罕见合影！
+        
+        12. 政治新闻
+        定义：聚焦政治事件、领导人动态、国际关系等内容，通常以新闻或分析的形式呈现，旨在揭示政治局势、政策变化或国际关系的动态。
+        标题示例：
+        中方外长行程有变，提前结束访欧匆匆回国，带回来一个好消息
+        宋庆龄在北京逝世后，远在美国的宋美龄只说了7个字，字字揪心！
+        庐山会议后，叶帅去劝彭德怀认个错，哭着说了一句心里话
+        
+        13. 历史人物
+        定义：聚焦于21世纪前具有重要影响的人物，包括他们的生平、事迹、成就、性格、趣事及其对历史进程的贡献。内容通常以传记、回忆录或历史分析的形式呈现，旨在还原人物的真实面貌并探讨其历史意义。
+        标题示例：
+        林彪去世后，蒋介石收到林彪与戴笠的一份密谈文件，看后拍桌大骂
+        张学良软禁时的一张实拍照片，头发秃顶，两眼无光，像个中年老头
+        1912年，孙中山和两个女儿罕见留影，面对镜头父女三人看起来很幸福
+        
+        14. 社会现象
+        定义：关注社会中出现的普遍现象、趋势或问题，通常涉及文化、经济、教育、民生等领域。内容以观察、分析或评论为主，旨在揭示现象背后的原因、影响及社会意义，引发公众的思考和讨论。
+        标题示例：
+        22年河南男子跳河救人，体力耗尽留遗言，被救女子猛然抓住他：一起走
+        浙江一老人刑满释放，靠蹬三轮为生，6年后，政府领导登门拜访：我们帮您分配工作
+        儿子收到清华通知书，父亲花5万请全村吃席，镇长看一眼竟说：这是假的
+        
+        15.财经科技
+        定义：聚焦于经济、金融、投资及行业发展的分析与预测，涵盖未来经济趋势、资产价值变化、行业变革及个人理财策略等内容。可以提供前瞻性的财经视角和实用的理财建议，帮助其把握经济动态、优化财务规划并应对行业变化。
+        标题示例：
+        未来10年，现金和房子都将贬值，只有2样东西最值钱
+        外卖时代将被终结？一个全新行业正悄悄取代外卖，你准备好了吗？
+        准备存款的一定要知道，今明两年，定期存款要记住“4不存”
+        
+        输入是一个 LIST， LIST 中的每个元素是一个元组，元组的第一个元素是文章的 ID，第二个元素是文章的标题。
+        最后输出结果请用JSON格式输出，key为ID，value为品类，仅输出JSON，不要markdown格式，不要任何其他内容
+        如果标题中包含半角双引号，则进行转义
+        输入的 LIST 是 {title_list}
+    """
+    return prompt
+
+
+
+

title_process_task.py

@@ -0,0 +1,16 @@
+"""
+@author: luojunhui
+"""
+from tasks.ai_tasks.category_generation_task import CategoryGenerationTask
+from tasks.title_rewrite_task import TitleRewriteTask
+
+
+if __name__ == '__main__':
+    # 1. 标题重写
+    title_rewrite_task = TitleRewriteTask()
+    title_rewrite_task.deal()
+
+    # 2. 标题分类
+    category_generation_task = CategoryGenerationTask()
+    category_generation_task.deal()
+