LongArticlesJob/tasks/title_rewrite_task.py

"""
@author: luojunhui
"""

import time
import traceback

from pymysql.cursors import DictCursor
from tqdm import tqdm

from applications import log
from applications.api import fetch_deepseek_response
from applications.const import TitleRewriteTaskConst
from applications.db import DatabaseConnector
from config import long_articles_config

const = TitleRewriteTaskConst()


def generate_prompt(ori_title):
    """
    生成prompt
    """
    prompt = f"""
    请将以下标题改写成适合公众号中小程序点击和传播的文章标题，文章标题的写作规范如下，请学习后进行文章标题的编写。直接输出最终的文章标题，文章标题撰写规范如下：
    1. 标题结构：要点前置，信息明确
        核心信息前置：标题开头直接点出文章的核心内容或亮点，吸引读者注意。例如：
          “我国存款最安全的五大银行，永远都不会倒闭，你知道是哪五家吗？”
          “亩产7000斤，被误认成萝卜却曾是‘救命粮’，如今成我国出口名蔬”。
        简洁明了：标题通常在20字以内，信息集中且易于理解。
        悬念前置结构：前半句设置反常/冲突场景（如"刑满释放蹬三轮"）+后半句用结果反转制造悬念（"政府领导登门分配工作"）
        多要素拼接：通过冒号/逗号分隔不同叙事主体（地域+人物冲突+权威评价），如"辽宁女子住高档小区被敲门，法院判决意外"
    
    2. 情绪表达：激发共鸣，引发好奇
        情感共鸣：通过情感化的语言触动读者，泪崩/守护/抱头痛哭等情感冲击词，配合家庭伦理场景
        例如：
          “老母亲分家产，给亲闺女30万，给养女一筐青菜，养女意外摔倒，看到筐子里的东西，瞬间愣住了”。
          “儿子卖车卖房给母亲治病，母亲去世后儿媳收拾房间，打开床底柜，儿子突然痛哭”。
        悬念与好奇心：通过提问或制造悬念，激发读者点击欲望。例如：
          “你知道是哪五家吗？”
          “打开床底柜，儿子突然痛哭”。
        冲突性情绪词：拍桌大骂/气愤不已/眼红不已/算计等强对抗性词汇
        结果反差刺激：用"风光善终/价值过亿/判决意外"等违反预期的结果
    
    3. 语言风格：口语化、接地气
        口语化表达：使用通俗易懂的语言，贴近读者生活。
        刻意使用"赶都赶不走/各吃各的/我就知道你在家"等市井化用语。
        例如：
          “狗屎运？江西男子钓鱼时发现青鱼尸骸，扒开后捡到鸡蛋大小的青鱼石”。
          “聪明的女人，不会帮婆家3种忙，而蠢女人才一再插手”。
        接地气的词汇：使用“狗屎运”“蠢女人”等口语化词汇，增强亲切感。
        身份反差构建：突出人物命运转折（老农→亿万富翁/囚犯→政府帮扶对象）
        权威背书暗示："专家气愤/法院判决/网友评价"等第三方视角增强可信度
    
    4. 标点运用：增强语气，突出重点
        问号与感叹号：通过问号制造悬念，感叹号强化情感。
        在关键转折点使用（"太气人了！/赔不了！"）
        问号制造互动：如"容嬷嬷是校花？"激发读者验证心理
        例如：
          “你知道是哪五家吗？”
          “太无耻了！湖南，一名厨师被公司派到云南‘出差’被拒……”
        引号与冒号：用于突出关键词或转折点。
        破折号递进：用"——"引导关键信息（"吃不完最好扔掉——"）
        例如：
          “被误认成萝卜却曾是‘救命粮’”。
          “女子归还后，失主拒绝支付报酬，还说：要有格局”。
    
    5. 热点与话题性：结合社会热点或争议
        社会热点：结合当前热点事件或争议话题，吸引关注。例如：
          “上海：男子超市连续购买46枚过期咸鸭蛋，2天分46次交易，向厂家索赔金14万，法院判了！”
        争议性话题：通过争议性内容引发讨论。例如：
          “李玉成终于说出实话，公开吐槽马玉琴年纪太大，结婚28年疑似后悔”。
    
    6. 数字与具体细节：增强可信度与吸引力
        数字的运用：通过具体数字增强标题的可信度和吸引力。例如：
          “亩产7000斤”。
          “22年河南男子跳河救人，体力耗尽留遗言”。
        细节描述：通过细节让标题更具画面感。例如：
          “打开床底柜，儿子突然痛哭”。
          “扒开后捡到鸡蛋大小的青鱼石”。
    
    7. 价值诉求：传递实用信息或情感价值
        实用信息：提供对读者有价值的信息。例如：
          “我国存款最安全的五大银行，永远都不会倒闭”。
          “72岁老人每天一个蒸苹果，半年后体检，看到指标变化让他乐开了花”。
        情感价值：通过情感故事或人生哲理打动读者。例如：
          “父母越老越能暴露家庭最真实的一面：当父母70岁，子女不该抱有这三种期待”。
    
    8. 名人效应与历史情怀：增强吸引力
        名人效应：提及名人或历史事件，吸引关注。例如：
          “难怪王扶林说陈晓旭不够漂亮，看看他选的原黛玉候选人，那才叫美”。
          “1975年‘下馆子’的老照片，2元能吃些什么，勾起那段最难忘的时光”。
    
    9.隐藏传播逻辑：通过标题中暗含的、能触发人性弱点（如猎奇、贪婪、同情）或社会痛点的心理机制，通过潜意识刺激读者点击欲望
       人性弱点触发：贪婪（200万保单）、猎奇（林彪密件）、窥私（家庭算计）
       生存焦虑关联：医疗（脑瘫儿）、养老（子女不孝）、食品安全（二次加热）
       身份代入设计：选择"老太太/外甥女/退休母亲"等易引发群体共鸣的角色
    输入的标题是： '{ori_title}'
    """
    return prompt


class TitleRewriteTask:
    """
    标题重写任务
    """

    def __init__(self):
        self.db = DatabaseConnector(db_config=long_articles_config)
        self.db.connect()

    def roll_back_blocked_tasks(self):
        """
        rollback blocked tasks
        """
        sql = f"""
            select id, title_rewrite_status_update_timestamp
            from publish_single_video_source
            where title_rewrite_status = {const.TITLE_REWRITE_LOCK_STATUS};
        """
        article_list = self.db.fetch(query=sql, cursor_type=DictCursor)
        if article_list:
            blocked_id_list = [
                i["id"]
                for i in article_list
                if (
                    int(time.time())
                    - i["title_rewrite_status_update_timestamp"]
                )
                > const.TITLE_REWRITE_LOCK_TIME
            ]
            if blocked_id_list:
                update_sql = f"""
                    update publish_single_video_source
                    set title_rewrite_status = %s
                    where id in %s and title_rewrite_status = %s;
                """
                self.db.save(
                    query=update_sql,
                    params=(
                        const.TITLE_REWRITE_INIT_STATUS,
                        tuple(blocked_id_list),
                        const.TITLE_REWRITE_LOCK_STATUS,
                    )
                )

    def get_articles_batch(self, batch_size=1000):
        """
        从数据库中获取文章
        """
        sql = f"""
            select content_trace_id, article_title
            from publish_single_video_source 
            where bad_status = {const.ARTICLE_POSITIVE_STATUS} 
                and audit_status = {const.ARTICLE_AUDIT_PASSED_STATUS} 
                and title_rewrite_status = {const.TITLE_REWRITE_INIT_STATUS}
                and platform in ('hksp', 'sph')
            limit {batch_size};
        """
        res = self.db.fetch(query=sql, cursor_type=DictCursor)
        return res

    def update_title_rewrite_status(self, content_trace_id, ori_status, new_status):
        """
        更新标题重写状态
        """
        sql = f"""
            update publish_single_video_source
            set title_rewrite_status = %s, title_rewrite_status_update_timestamp = %s
            where content_trace_id = %s and title_rewrite_status= %s;
        """
        affected_rows = self.db.save(
            query=sql, params=(new_status, int(time.time()), content_trace_id, ori_status)
        )
        return affected_rows

    def insert_into_rewrite_table(self, content_trace_id, new_title):
        """
        insert into rewrite_table
        """
        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,
                const.PROMPT_VERSION
            ),
        )

    def rewrite_each_article(self, article):
        """
        rewrite each article
        """
        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,
        )
        if not affected_rows:
            return

        try:
            prompt = generate_prompt(article_title)
            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
            )

            # 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,
            )
        except Exception as e:
            log(
                task="title rewrite task",
                function="rewrite_each_article",
                message=content_trace_id,
                status="fail",
                data={
                    "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,
            )

    def deal(self):
        """
        get tasks && deal tasks
        """
        # rollback blocked tasks
        try:
            self.roll_back_blocked_tasks()
        except Exception as e:
            log(
                task="title rewrite task",
                function="roll_back_blocked_tasks",
                message="roll back blocked tasks fail",
                status="fail",
                data={
                    "error_message": str(e),
                    "error_type": type(e).__name__,
                    "traceback": traceback.format_exc()
                }
            )

        # process tasks
        articles = self.get_articles_batch()
        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"]})