LongArticleTaskServer/applications/tasks/cold_start_tasks/article_pool_cold_start.py

from __future__ import annotations

import time
import asyncio
import datetime
import traceback

from typing import List
from pandas import DataFrame
from tqdm.asyncio import tqdm

from applications.api import task_apollo, feishu_robot
from applications.api import auto_create_crawler_task
from applications.api import auto_bind_crawler_task_to_generate_task
from applications.config import cold_start_category_map, input_source_map
from applications.utils import get_titles_from_produce_plan
from applications.tasks.cold_start_tasks.article_pool import (
    ArticlePoolColdStartStrategy,
    ArticlePoolFilterStrategy
)


class ArticlePoolColdStart(ArticlePoolColdStartStrategy, ArticlePoolFilterStrategy):
    def __init__(self, pool, log_client, trace_id):
        super().__init__(pool, log_client, trace_id)

    async def get_article_from_meta_table(
        self, platform: str, crawl_method: str, strategy: str, category: str | None
    ) -> DataFrame:
        """
        @param platform: 文章抓取平台
        @param crawl_method: 文章抓取模式
        @param strategy: 供给策略
        """
        match platform:
            case "weixin":
                article_list = await self.get_weixin_cold_start_articles(
                    crawl_method, strategy, category
                )
            case "toutiao":
                article_list = await self.get_toutiao_cold_start_articles(
                    crawl_method, strategy, category
                )
            case _:
                raise ValueError("Invalid platform")
        return DataFrame(
            article_list,
            columns=[
                "article_id",
                "title",
                "link",
                "llm_sensitivity",
                "score",
                "category_by_ai",
            ],
        )

    async def filter_published_titles(self, plan_id):
        """
        过滤已添加至aigc中的标题
        """
        published_title_tuple = await get_titles_from_produce_plan(self.pool, plan_id)
        if not published_title_tuple:
            return 0
        update_query = """
            update crawler_meta_article set status = %s where title in %s and status = %s;
        """
        changed_rows = await self.pool.async_save(
            query=update_query,
            params=(self.PUBLISHED_STATUS, published_title_tuple, self.INIT_STATUS),
        )
        return changed_rows

    async def insert_crawler_plan_into_database(
        self, crawler_plan_id, crawler_plan_name, create_timestamp
    ):
        query = """
            insert into article_crawler_plan (crawler_plan_id, name, create_timestamp)
            values (%s, %s, %s)
        """
        try:
            await self.pool.async_save(
                query=query,
                params=(crawler_plan_id, crawler_plan_name, create_timestamp),
            )
        except Exception as e:
            await feishu_robot.bot(
                title="品类冷启任务，记录抓取计划id失败",
                detail={
                    "error": str(e),
                    "error_msg": traceback.format_exc(),
                    "crawler_plan_id": crawler_plan_id,
                    "crawler_plan_name": crawler_plan_name,
                },
            )

    async def change_article_status_while_publishing(self, article_id_list):
        """
        :param: article_id_list: 文章的唯一 id
        :return:
        """
        query = """
            update crawler_meta_article
            set status = %s
            where article_id in %s and status = %s;
        """
        affect_rows = await self.pool.async_save(
            query=query,
            params=(self.PUBLISHED_STATUS, tuple(article_id_list), self.INIT_STATUS),
        )
        return affect_rows

    async def create_crawler_plan_and_bind_to_produce_plan(
            self,
            strategy: str,
            crawl_method: str,
            category: str,
            platform: str,
            url_list: List[str],
            plan_id: str,
    ):
        # create_crawler_plan
        crawler_plan_response = await auto_create_crawler_task(
            plan_id=None,
            plan_name=f"冷启动-{strategy}-{category}-{datetime.date.today().__str__()}-{len(url_list)}",
            plan_tag="品类冷启动",
            platform=platform,
            url_list=url_list,
        )
        # save to db
        create_timestamp = int(time.time()) * 1000
        crawler_plan_id = crawler_plan_response["data"]["id"]
        crawler_plan_name = crawler_plan_response["data"]["name"]
        await self.insert_crawler_plan_into_database(
            crawler_plan_id, crawler_plan_name, create_timestamp
        )

        # auto bind to generate plan
        new_crawler_task_list = [
            {
                "contentType": 1,
                "inputSourceType": 2,
                "inputSourceSubType": None,
                "fieldName": None,
                "inputSourceValue": crawler_plan_id,
                "inputSourceLabel": crawler_plan_name,
                "inputSourceModal": 3,
                "inputSourceChannel": input_source_map[platform],
            }
        ]
        generate_plan_response = await auto_bind_crawler_task_to_generate_task(
            crawler_task_list=new_crawler_task_list, generate_task_id=plan_id
        )
        await self.log_client.log(
            contents={
                "task": "article_pool_cold_start",
                "platform": platform,
                "crawl_method": crawl_method,
                "status": "success",
                "trace_id": self.trace_id,
                "message": "绑定至生成计划成功",
                "data": generate_plan_response,
            }
        )

    async def create_cold_start_plan(
        self,
        platform,
        plan_id,
        strategy="strategy_v1",
        category=None,
        crawl_method=None,
    ):
        # get article data_frame from meta article
        article_dataframe = await self.get_article_from_meta_table(
            platform, crawl_method, strategy, category
        )

        await self.log_client.log(
            contents={
                "task": "article_pool_cold_start",
                "platform": platform,
                "crawl_method": crawl_method,
                "status": "success",
                "trace_id": self.trace_id,
                "message": "获取文章成功",
                "data": {"article_length": article_dataframe.shape[0]},
            }
        )

        filter_article_df = await self.article_pool_filter(
            strategy, platform, article_dataframe, crawl_method, category
        )
        match strategy:
            case "strategy_v1":
                # split article into each category
                category_list = await task_apollo.get_config_value(key="category_list")
                for ai_category in category_list:
                    filter_category_df = filter_article_df[
                        filter_article_df["category_by_ai"] == ai_category
                    ]
                    url_list = filter_category_df["link"].values.tolist()
                    if url_list:
                        await self.create_crawler_plan_and_bind_to_produce_plan(
                            strategy, crawl_method, ai_category, platform, url_list, plan_id
                        )
                        # change article status
                        article_id_list = filter_category_df["article_id"].values.tolist()
                        await self.change_article_status_while_publishing(
                            article_id_list=article_id_list
                        )

            case "strategy_v2":
                url_list =  filter_article_df["link"].values.tolist()
                await self.create_crawler_plan_and_bind_to_produce_plan(
                    strategy, crawl_method, category, platform, url_list, plan_id
                )
                # change article status
                article_id_list = filter_article_df["article_id"].values.tolist()
                await self.change_article_status_while_publishing(
                    article_id_list=article_id_list
                )


    async def deal(
        self,
        platform: str,
        strategy="strategy_v1",
        crawl_methods=None,
        category_list=None,
    ) -> None:
        """execute cold start task in different strategy"""
        match strategy:
            case "strategy_v1":
                if not crawl_methods:
                    crawl_methods = self.DEFAULT_CRAWLER_METHODS

                await self.log_client.log(
                    contents={
                        "task": "article_pool_cold_start",
                        "platform": platform,
                        "crawl_methods": crawl_methods,
                        "status": "success",
                        "trace_id": self.trace_id,
                    }
                )

                crawl_methods_map = await task_apollo.get_config_value(
                    key="category_cold_start_map"
                )

                for crawl_method in crawl_methods:
                    try:
                        plan_id = crawl_methods_map[crawl_method]
                        affected_rows = await self.filter_published_titles(plan_id)
                        await self.log_client.log(
                            contents={
                                "task": "article_pool_cold_start",
                                "platform": platform,
                                "crawl_method": crawl_method,
                                "status": "success",
                                "trace_id": self.trace_id,
                                "message": "通过已抓取标题修改文章状态",
                                "data": {"affected_rows": affected_rows},
                            }
                        )
                        await self.create_cold_start_plan(
                            platform=platform,
                            plan_id=plan_id,
                            crawl_method=crawl_method,
                        )

                    except Exception as e:
                        await feishu_robot.bot(
                            title="文章冷启动异常",
                            detail={
                                "crawl_method": crawl_method,
                                "error": str(e),
                                "function": "deal",
                                "traceback": traceback.format_exc(),
                            },
                        )

            case "strategy_v2":
                if not category_list:
                    category_list = list(cold_start_category_map.keys())

                for category in tqdm(category_list):
                    try:
                        plan_id = cold_start_category_map[category]
                        affected_rows = await self.filter_published_titles(plan_id)
                        await self.log_client.log(
                            contents={
                                "task": "article_pool_cold_start",
                                "platform": platform,
                                "category": category,
                                "status": "success",
                                "trace_id": self.trace_id,
                                "message": "通过已抓取标题修改文章状态",
                                "data": {"affected_rows": affected_rows},
                            }
                        )
                        await self.create_cold_start_plan(
                            platform=platform,
                            strategy=strategy,
                            plan_id=plan_id,
                            category=category,
                        )
                        await asyncio.sleep(120)
                    except Exception as e:
                        await feishu_robot.bot(
                            title="文章冷启动异常",
                            detail={
                                "category": category,
                                "strategy": strategy,
                                "error": str(e),
                                "function": "deal",
                                "traceback": traceback.format_exc(),
                            },
                        )

                if self.cold_start_records:
                    columns = [
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="plain_text",
                            sheet_name="category",
                            display_name="文章品类",
                        ),
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="number",
                            sheet_name="cold_start_num",
                            display_name="本次冷启数量",
                        ),
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="number",
                            sheet_name="total_length",
                            display_name="总文章剩余数量",
                        ),
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="number",
                            sheet_name="filter_by_title_length",
                            display_name="标题长度过滤",
                        ),
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="number",
                            sheet_name="filter_by_sensitivity",
                            display_name="敏感词过滤",
                        ),
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="number",
                            sheet_name="filter_by_llm_sensitity",
                            display_name="经过大模型判断敏感过滤",
                        ),
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="number",
                            sheet_name="filter_by_score",
                            display_name="经过相关性分过滤",
                        ),
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="number",
                            sheet_name="read_avg_threshold",
                            display_name="阅读均值倍数阈值",
                        ),
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="number",
                            sheet_name="read_threshold",
                            display_name="阅读量阈值",
                        ),
                        feishu_robot.create_feishu_columns_sheet(
                            sheet_type="number",
                            sheet_name="title_length_threshold",
                            display_name="标题长度阈值",
                        ),
                    ]
                    await feishu_robot.bot(
                        title="长文文章路冷启动发布",
                        detail={
                            "columns": columns,
                            "rows": self.cold_start_records,
                        },
                        table=True,
                        mention=False,
                    )

            case _:
                raise Exception(f"error strategy {strategy}")