#!/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.custom_keywords_path: Optional[Path] = self._resolve_custom_keywords_path() 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 getattr(settings, "ENABLE_KEYWORD_SEARCH", True): logger.info("关键词搜索已通过配置禁用，返回默认关键词列表") return self._limit_keywords(self._get_default_keywords(), max_keywords) if not target_date: target_date = date.today() if getattr(settings, "USE_DEFAULT_KEYWORDS_ONLY", False): logger.info("配置启用默认关键词模式，直接返回默认关键词") return self._limit_keywords(self._get_default_keywords(), max_keywords) logger.info(f"正在获取 {target_date} 的关键词...") # 优先使用自定义关键词 custom_keywords = self._get_custom_keywords(target_date, max_keywords) if custom_keywords: return custom_keywords # 首先尝试获取指定日期的关键词 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)} 个关键词") return self._limit_keywords(keywords, max_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)) limited_keywords = self._limit_keywords(unique_keywords, max_keywords) logger.info(f"从最近7天的数据中获取到 {len(limited_keywords)} 个关键词") return limited_keywords # 如果都没有，返回默认关键词 logger.info("没有找到任何关键词数据，使用默认关键词") return self._limit_keywords(self._get_default_keywords(), max_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 _resolve_custom_keywords_path(self) -> Optional[Path]: """解析自定义关键词文件路径""" raw_path = getattr(settings, "CUSTOM_KEYWORDS_FILE", None) if not raw_path: return None path = Path(raw_path).expanduser() if not path.is_absolute(): path = project_root / path return path @staticmethod def _limit_keywords(keywords: List[str], max_keywords: int) -> List[str]: """根据最大数量限制关键词""" if not keywords: return [] if max_keywords and len(keywords) > max_keywords: keywords = random.sample(keywords, max_keywords) return keywords def _get_custom_keywords(self, target_date: date, max_keywords: int) -> Optional[List[str]]: """ 从自定义关键词文件中获取指定日期的关键词 支持格式： 1. ["AI", "大模型"] 2. {"2025-11-26": ["AI"], "default": ["科技"]} """ if not self.custom_keywords_path: return None path = self.custom_keywords_path if not path.exists(): logger.warning(f"自定义关键词文件不存在: {path}") return None try: with path.open("r", encoding="utf-8") as f: data = json.load(f) except Exception as e: logger.error(f"读取自定义关键词文件失败({path}): {e}") return None keywords: Optional[List[str]] = None if isinstance(data, list): keywords = data elif isinstance(data, dict): date_key = target_date.isoformat() if date_key in data and isinstance(data[date_key], list): keywords = data[date_key] elif "default" in data and isinstance(data["default"], list): keywords = data["default"] elif "keywords" in data and isinstance(data["keywords"], list): keywords = data["keywords"] if not keywords: return None cleaned_keywords = [kw.strip() for kw in keywords if isinstance(kw, str) and kw.strip()] if not cleaned_keywords: return None limited = self._limit_keywords(cleaned_keywords, max_keywords) logger.info(f"使用自定义关键词（{path}）共 {len(limited)} 个") return limited def _get_default_keywords(self) -> List[str]: """获取默认关键词列表""" return [ "F6智慧门店","F6智数","中国汽车后市场白皮书","南京爱福路汽车科技有限公司","汽车后市场","汽车修理厂", "新康众","天猫养车","汽后","汽修厂","爱福路","康众", ] 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() # 当配置为仅使用默认关键词时，直接返回默认关键词的摘要，避免上层误判为“无数据” if getattr(settings, "USE_DEFAULT_KEYWORDS_ONLY", False): default_keywords = self._get_default_keywords() return { 'date': target_date, 'keywords_count': len(default_keywords), 'summary': '使用默认关键词模式', 'has_data': bool(default_keywords) } 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("关键词管理器测试完成！")