nanocode/xhs_trending_automation/scripts/workflow_from_it_db.py

#!/usr/bin/env python3
"""
小红书工作流 - 从IT资讯自动化数据库获取数据源

与IT资讯自动化流程整合：
- IT资讯自动化：抓取 → 翻译/总结 → 保存到数据库 → 发邮件
- 小红书工作流：从同一数据库读取 → 挑选文章 → 发布小红书
"""
import os
import sys
import json
import sqlite3
import argparse
from datetime import datetime
from typing import Optional, Dict, List

# 导入发布模块
sys.path.insert(0, os.path.dirname(__file__))
from publisher_adapter import XHSPublisherAdapter


class XHSFromITDBWorkflow:
    """
    小红书工作流 - IT资讯数据库版
    
    数据源: /home/ubuntu/nanocode/it_automation_enhanced/data/articles.db
    """
    
    # IT资讯数据库路径
    IT_DB_PATH = "/home/ubuntu/nanocode/it_automation_enhanced/data/articles.db"
    IT_OUTPUT_DIR = "/home/ubuntu/nanocode/it_automation_enhanced/output"
    
    def __init__(self, config_path: str = None):
        self.config = self._load_config(config_path)
        self.publisher = XHSPublisherAdapter()
        
        # 路径
        self.log_dir = self.config.get('paths', {}).get('log_dir', './logs')
        self.output_dir = self.config.get('paths', {}).get('output_dir', './output')
        os.makedirs(self.log_dir, exist_ok=True)
        os.makedirs(self.output_dir, exist_ok=True)
    
    def _load_config(self, config_path: str) -> dict:
        """加载配置"""
        if not config_path:
            config_path = os.path.join(
                os.path.dirname(__file__), '..', 'config', 'config.json'
            )
        
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"配置加载失败: {e}，使用默认配置")
            return {
                'workflow': {'content_style': 'tech'},
                'paths': {'log_dir': './logs', 'output_dir': './output'}
            }
    
    def _log(self, message: str):
        """记录日志"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        log_message = f"[{timestamp}] {message}"
        print(log_message)
        
        log_file = os.path.join(
            self.log_dir, 
            f"xhs_from_itdb_{datetime.now().strftime('%Y%m%d')}.log"
        )
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(log_message + '\n')
    
    def get_processed_articles(self, limit: int = 10) -> List[Dict]:
        """从IT资讯数据库获取已处理但未发小红书的文章"""
        articles = []
        
        if not os.path.exists(self.IT_DB_PATH):
            self._log(f"❌ IT资讯数据库不存在: {self.IT_DB_PATH}")
            return articles
        
        try:
            conn = sqlite3.connect(self.IT_DB_PATH)
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            
            # 查询已处理但未发小红书的文章
            cursor.execute("""
                SELECT id, title, url, source, category, 
                       publish_date, word_count, created_at, xhs_status
                FROM articles 
                WHERE processed = 1 
                  AND (xhs_status = 0 OR xhs_status IS NULL)
                ORDER BY created_at DESC
                LIMIT ?
            """, (limit,))
            
            rows = cursor.fetchall()
            for row in rows:
                articles.append({
                    'id': row['id'],
                    'title': row['title'],
                    'url': row['url'],
                    'source': row['source'],
                    'category': row['category'],
                    'publish_date': row['publish_date'],
                    'word_count': row['word_count'],
                    'created_at': row['created_at'],
                    'xhs_status': row['xhs_status'] or 0
                })
            
            conn.close()
            self._log(f"📊 找到 {len(articles)} 篇已处理但未发小红书的文章")
            
        except Exception as e:
            self._log(f"❌ 数据库查询失败: {e}")
        
        return articles
    
    def load_processed_content(self, article_id: int) -> Optional[Dict]:
        """加载IT资讯自动化的处理结果"""
        processed_file = os.path.join(
            self.IT_OUTPUT_DIR, 
            f"article_{article_id}_processed.json"
        )
        
        if not os.path.exists(processed_file):
            self._log(f"⚠️ 处理结果文件不存在: {processed_file}")
            return None
        
        try:
            with open(processed_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 构建小红书内容
            content = {
                'article_id': article_id,
                'title': data.get('translated_title', data.get('original_title', '')),
                'body': data.get('summary', data.get('translated_content', ''))[:900],
                'tags': self._generate_tags(data),
                'source': data.get('source', '技术'),
                'url': data.get('url', ''),
                'original_title': data.get('original_title', ''),
            }
            
            return content
            
        except Exception as e:
            self._log(f"❌ 加载处理结果失败: {e}")
            return None
    
    def _generate_tags(self, data: Dict) -> List[str]:
        """根据文章数据生成标签"""
        source = data.get('source', '技术')
        category = data.get('category', '科技')
        
        # 基础标签
        base_tags = ['科技资讯', '技术分享']
        
        # 根据来源添加标签
        source_tags = {
            'InfoQ': 'InfoQ',
            'Wired': 'Wired',
            'TechCrunch': 'TechCrunch',
        }
        
        if source in source_tags:
            base_tags.append(source_tags[source])
        
        # 根据分类添加标签
        if category and category not in base_tags:
            base_tags.append(category)
        
        # 确保5个标签
        while len(base_tags) < 5:
            base_tags.append('程序员')
        
        return base_tags[:5]
    
    def mark_as_published(self, article_id: int) -> bool:
        """标记文章为已发布到小红书"""
        try:
            conn = sqlite3.connect(self.IT_DB_PATH)
            cursor = conn.cursor()
            
            cursor.execute("""
                UPDATE articles 
                SET xhs_status = 1, updated_at = datetime('now')
                WHERE id = ?
            """, (article_id,))
            
            conn.commit()
            conn.close()
            
            self._log(f"✅ 文章 {article_id} 已标记为已发布")
            return True
            
        except Exception as e:
            self._log(f"❌ 标记发布状态失败: {e}")
            return False
    
    def run(self, article_id: int = None, dry_run: bool = False) -> bool:
        """
        执行小红书发布工作流
        
        Args:
            article_id: 指定文章ID（可选，不指定则自动选择）
            dry_run: 试运行模式
            
        Returns:
            是否成功
        """
        self._log("="*60)
        self._log("🚀 小红书工作流启动 (IT资讯数据源)")
        self._log("="*60)
        
        # 步骤1: 获取文章
        if article_id:
            # 指定ID
            self._log(f"\n📄 步骤1: 使用指定文章 ID={article_id}")
            articles = self.get_processed_articles(limit=50)
            article = next((a for a in articles if a['id'] == article_id), None)
            
            if not article:
                # 尝试从数据库直接查询
                self._log(f"⚠️ 文章 {article_id} 不在未发布列表中，尝试直接加载...")
                content = self.load_processed_content(article_id)
                if content:
                    content['article_id'] = article_id
                else:
                    self._log(f"❌ 无法加载文章 {article_id}")
                    return False
            else:
                content = self.load_processed_content(article_id)
        else:
            # 自动选择
            self._log("\n📄 步骤1: 从IT资讯数据库获取文章...")
            articles = self.get_processed_articles(limit=5)
            
            if not articles:
                self._log("❌ 没有可发布的文章（IT资讯可能尚未运行或文章都已发布）")
                self._log("💡 提示: IT资讯自动化每天16:00运行，请等待或手动运行")
                return False
            
            # 选择第一篇
            article = articles[0]
            self._log(f"✅ 选中文章: {article['title'][:50]}...")
            self._log(f"   来源: {article['source']} | 字数: {article['word_count']}")
            
            content = self.load_processed_content(article['id'])
        
        if not content:
            self._log("❌ 无法加载文章内容")
            return False
        
        # 步骤2: 显示内容预览
        self._log("\n📝 步骤2: 内容预览")
        self._log(f"标题: {content['title'][:60]}...")
        self._log(f"正文: {len(content['body'])} 字符")
        self._log(f"标签: {', '.join(content['tags'])}")
        
        # 保存到本地
        output_file = os.path.join(
            self.output_dir,
            f"xhs_content_{content['article_id']}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        )
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(content, f, ensure_ascii=False, indent=2)
        self._log(f"\n💾 内容已保存: {output_file}")
        
        # 步骤3: 发布
        if dry_run:
            self._log("\n🧪 试运行模式，跳过发布")
            self._log("="*60)
            return True
        
        self._log("\n📱 步骤3: 发布到小红书...")
        success = self.publisher.publish(
            title=content['title'],
            content=content['body'],
            tags=content['tags'],
            style=self.config.get('workflow', {}).get('content_style', 'tech')
        )
        
        if success:
            self._log("🎉 发布成功！")
            # 标记为已发布
            self.mark_as_published(content['article_id'])
        else:
            self._log("❌ 发布失败")
        
        self._log("="*60)
        return success


def main():
    parser = argparse.ArgumentParser(description='小红书工作流 - IT资讯数据源')
    parser.add_argument('--article-id', type=int, help='指定文章ID')
    parser.add_argument('--dry-run', action='store_true', help='试运行不发布')
    parser.add_argument('--config', help='配置文件路径')
    args = parser.parse_args()
    
    workflow = XHSFromITDBWorkflow(config_path=args.config)
    success = workflow.run(
        article_id=args.article_id,
        dry_run=args.dry_run
    )
    
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()