bettafish-company/MindSpider/DeepSentimentCrawling/keyword_manager.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DeepSentimentCrawling模块 - 关键词管理器
从BroadTopicExtraction模块获取关键词并分配给不同平台进行爬取
"""

import sys
import json
from datetime import date, timedelta, datetime
from pathlib import Path
from typing import List, Dict, Optional
import random
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root))

try:
    import config
except ImportError:
    raise ImportError("无法导入config.py配置文件")

from config import settings
from loguru import logger

class KeywordManager:
    """关键词管理器"""
    
    def __init__(self):
        """初始化关键词管理器"""
        self.engine: Engine = None
        self.connect()
    
    def connect(self):
        """连接数据库"""
        try:
            dialect = (settings.DB_DIALECT or "mysql").lower()
            if dialect in ("postgresql", "postgres"):
                url = f"postgresql+psycopg://{settings.DB_USER}:{settings.DB_PASSWORD}@{settings.DB_HOST}:{settings.DB_PORT}/{settings.DB_NAME}"
            else:
                url = f"mysql+pymysql://{settings.DB_USER}:{settings.DB_PASSWORD}@{settings.DB_HOST}:{settings.DB_PORT}/{settings.DB_NAME}?charset={settings.DB_CHARSET}"
            self.engine = create_engine(url, future=True)
            logger.info(f"关键词管理器成功连接到数据库: {settings.DB_NAME}")
        except ModuleNotFoundError as e:
            missing: str = str(e)
            if "psycopg" in missing:
                logger.error("数据库连接失败: 未安装PostgreSQL驱动 psycopg。请安装: psycopg[binary]。参考指令：uv pip install psycopg[binary]")
            elif "pymysql" in missing:
                logger.error("数据库连接失败: 未安装MySQL驱动 pymysql。请安装: pymysql。参考指令：uv pip install pymysql")
            else:
                logger.error(f"数据库连接失败(缺少驱动): {e}")
            raise
        except Exception as e:
            logger.exception(f"关键词管理器数据库连接失败: {e}")
            raise
    
    def get_latest_keywords(self, target_date: date = None, max_keywords: int = 100) -> List[str]:
        """
        获取最新的关键词列表
        
        Args:
            target_date: 目标日期，默认为今天
            max_keywords: 最大关键词数量
        
        Returns:
            关键词列表
        """
        if not target_date:
            target_date = date.today()
        
        logger.info(f"正在获取 {target_date} 的关键词...")
        
        # 首先尝试获取指定日期的关键词
        topics_data = self.get_daily_topics(target_date)
        
        if topics_data and topics_data.get('keywords'):
            keywords = topics_data['keywords']
            logger.info(f"成功获取 {target_date} 的 {len(keywords)} 个关键词")
            
            # 如果关键词太多，随机选择指定数量
            if len(keywords) > max_keywords:
                keywords = random.sample(keywords, max_keywords)
                logger.info(f"随机选择了 {max_keywords} 个关键词")
            
            return keywords
        
        # 如果没有当天的关键词，尝试获取最近几天的
        logger.info(f"{target_date} 没有关键词数据，尝试获取最近的关键词...")
        recent_topics = self.get_recent_topics(days=7)
        
        if recent_topics:
            # 合并最近几天的关键词
            all_keywords = []
            for topic in recent_topics:
                if topic.get('keywords'):
                    all_keywords.extend(topic['keywords'])
            
            # 去重并限制数量
            unique_keywords = list(set(all_keywords))
            if len(unique_keywords) > max_keywords:
                unique_keywords = random.sample(unique_keywords, max_keywords)
            
            logger.info(f"从最近7天的数据中获取到 {len(unique_keywords)} 个关键词")
            return unique_keywords
        
        # 如果都没有，返回默认关键词
        logger.info("没有找到任何关键词数据，使用默认关键词")
        return self._get_default_keywords()
    
    def get_daily_topics(self, extract_date: date = None) -> Optional[Dict]:
        """
        获取每日话题分析
        
        Args:
            extract_date: 提取日期，默认为今天
        
        Returns:
            话题分析数据，如果不存在返回None
        """
        if not extract_date:
            extract_date = date.today()
        
        try:
            with self.engine.connect() as conn:
                result = conn.execute(
                    text("SELECT * FROM daily_topics WHERE extract_date = :d"),
                    {"d": extract_date},
                ).mappings().first()
            
            if result:
                # 转为可变dict再赋值
                result = dict(result)
                result['keywords'] = json.loads(result['keywords']) if result.get('keywords') else []
                return result
            else:
                return None
                
        except Exception as e:
            logger.exception(f"获取话题分析失败: {e}")
            return None
    
    def get_recent_topics(self, days: int = 7) -> List[Dict]:
        """
        获取最近几天的话题分析
        
        Args:
            days: 天数
        
        Returns:
            话题分析列表
        """
        try:
            start_date = date.today() - timedelta(days=days)
            with self.engine.connect() as conn:
                results = conn.execute(
                    text(
                        """
                        SELECT * FROM daily_topics 
                        WHERE extract_date >= :start_date
                        ORDER BY extract_date DESC
                        """
                    ),
                    {"start_date": start_date},
                ).mappings().all()
            
            # 转为可变dict列表再处理
            results = [dict(r) for r in results]
            for result in results:
                result['keywords'] = json.loads(result['keywords']) if result.get('keywords') else []
            
            return results
            
        except Exception as e:
            logger.exception(f"获取最近话题分析失败: {e}")
            return []
    
    def _get_default_keywords(self) -> List[str]:
        """获取默认关键词列表"""
        return [
            "科技", "人工智能", "AI", "编程", "互联网",
            "创业", "投资", "理财", "股市", "经济",
            "教育", "学习", "考试", "大学", "就业",
            "健康", "养生", "运动", "美食", "旅游",
            "时尚", "美妆", "购物", "生活", "家居",
            "电影", "音乐", "游戏", "娱乐", "明星",
            "新闻", "热点", "社会", "政策", "环保"
        ]
    
    def get_all_keywords_for_platforms(self, platforms: List[str], target_date: date = None, 
                                      max_keywords: int = 100) -> List[str]:
        """
        为所有平台获取相同的关键词列表
        
        Args:
            platforms: 平台列表
            target_date: 目标日期
            max_keywords: 最大关键词数量
        
        Returns:
            关键词列表（所有平台共用）
        """
        keywords = self.get_latest_keywords(target_date, max_keywords)
        
        if keywords:
            logger.info(f"为 {len(platforms)} 个平台准备了相同的 {len(keywords)} 个关键词")
            logger.info(f"每个关键词将在所有平台上进行爬取")
        
        return keywords
    
    def get_keywords_for_platform(self, platform: str, target_date: date = None, 
                                max_keywords: int = 50) -> List[str]:
        """
        为特定平台获取关键词（现在所有平台使用相同关键词）
        
        Args:
            platform: 平台名称
            target_date: 目标日期
            max_keywords: 最大关键词数量
        
        Returns:
            关键词列表（与其他平台相同）
        """
        keywords = self.get_latest_keywords(target_date, max_keywords)
        
        logger.info(f"为平台 {platform} 准备了 {len(keywords)} 个关键词（与其他平台相同）")
        return keywords
    
    def _filter_keywords_by_platform(self, keywords: List[str], platform: str) -> List[str]:
        """
        根据平台特性过滤关键词
        
        Args:
            keywords: 原始关键词列表
            platform: 平台名称
        
        Returns:
            过滤后的关键词列表
        """
        # 平台特性关键词映射（可以根据需要调整）
        platform_preferences = {
            'xhs': ['美妆', '时尚', '生活', '美食', '旅游', '购物', '健康', '养生'],
            'dy': ['娱乐', '音乐', '舞蹈', '搞笑', '美食', '生活', '科技', '教育'],
            'ks': ['生活', '搞笑', '农村', '美食', '手工', '音乐', '娱乐'],
            'bili': ['科技', '游戏', '动漫', '学习', '编程', '数码', '科普'],
            'wb': ['热点', '新闻', '娱乐', '明星', '社会', '时事', '科技'],
            'tieba': ['游戏', '动漫', '学习', '生活', '兴趣', '讨论'],
            'zhihu': ['知识', '学习', '科技', '职场', '投资', '教育', '思考']
        }
        
        # 如果平台有特定偏好，优先选择相关关键词
        preferred_keywords = platform_preferences.get(platform, [])
        
        if preferred_keywords:
            # 先选择平台偏好的关键词
            filtered = []
            remaining = []
            
            for keyword in keywords:
                if any(pref in keyword for pref in preferred_keywords):
                    filtered.append(keyword)
                else:
                    remaining.append(keyword)
            
            # 如果偏好关键词不够，补充其他关键词
            if len(filtered) < len(keywords) // 2:
                filtered.extend(remaining[:len(keywords) - len(filtered)])
            
            return filtered
        
        # 如果没有特定偏好，返回原关键词
        return keywords
    
    def get_crawling_summary(self, target_date: date = None) -> Dict:
        """
        获取爬取任务摘要
        
        Args:
            target_date: 目标日期
        
        Returns:
            爬取摘要信息
        """
        if not target_date:
            target_date = date.today()
        
        topics_data = self.get_daily_topics(target_date)
        
        if topics_data:
            return {
                'date': target_date,
                'keywords_count': len(topics_data.get('keywords', [])),
                'summary': topics_data.get('summary', ''),
                'has_data': True
            }
        else:
            return {
                'date': target_date,
                'keywords_count': 0,
                'summary': '暂无数据',
                'has_data': False
            }
    
    def close(self):
        """关闭数据库连接"""
        if self.engine:
            self.engine.dispose()
            logger.info("关键词管理器数据库连接已关闭")
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

if __name__ == "__main__":
    # 测试关键词管理器
    with KeywordManager() as km:
        # 测试获取关键词
        keywords = km.get_latest_keywords(max_keywords=20)
        logger.info(f"获取到的关键词: {keywords}")
        
        # 测试平台分配
        platforms = ['xhs', 'dy', 'bili']
        distribution = km.distribute_keywords_by_platform(keywords, platforms)
        for platform, kws in distribution.items():
            logger.info(f"{platform}: {kws}")
        
        # 测试爬取摘要
        summary = km.get_crawling_summary()
        logger.info(f"爬取摘要: {summary}")
        
        logger.info("关键词管理器测试完成！")