meme/scheduler/jobs.py

import json
from llm import prompt
from utils.logger import logger
import datetime
from llm.generate_podcast import generate_topics
from models.script import Script
from config.database import SessionLocal


def job_heartbeat():
    logger.info(f"[heartbeat] {datetime.datetime.now()}")

def job_generate_topics():
    """定时任务：搜索上一周的热门梗并保存至数据库。"""

    # 1. 调用 LLM 生成热门梗列表
    topics = generate_topics()
    content = {"topics": topics}
    if not topics:
        logger.warning("No topics generated.")
        return
    
    # 2. 构建 Script 实例
    # subject 以当前日期为准，格式 YYYY-MM-DD
    today_str = datetime.datetime.now().strftime("%Y-%m-%d")
    db = SessionLocal()
    try:
        # 查询是否已存在 project+subject 唯一记录
        script = db.query(Script).filter_by(project="梗文化研究所", subject=today_str).first()
        if script:
            # 存在则更新内容
            script.content = json.dumps(content, ensure_ascii=False, indent=2)
            db.commit()
            logger.info(f"Updated script for {today_str} with {len(topics)} topics.")
        else:
            # 不存在则新建
            script = Script(
                project="梗文化研究所",
                subject=today_str,
                content=json.dumps(content, ensure_ascii=False, indent=2)
            )
            db.add(script)
            db.commit()
            logger.info(f"Saved script for {today_str} with {len(topics)} topics.")
    except Exception as e:
        db.rollback()
        logger.error(f"Failed to save/update script for {today_str}: {e}")


def job_generate_bits():
    """定时任务：为最新梗生成脱口秀段子并保存至数据库。"""
    db = SessionLocal()
    try:
        # 获取最新的 Script 记录
        script = db.query(Script).filter_by(project="梗文化研究所").order_by(Script.create_time.desc()).first()
        if not script or not script.content:
            logger.warning("No script found for generating bits.")
            return
        
        data = json.loads(script.content)
        topics = data.get("topics", [])
        if not topics:
            logger.warning("No topics in the latest script.")
            return
        
        # 仅处理第一个梗
        top = topics[0]
        meme_name = top.get("title") or top.get("name") or "未知梗"
        
        # 构建研究文本
        parts = []
        if "summary" in top:
            parts.append(f"简介：{top['summary']}")
        if "origin" in top:
            parts.append(f"可能起源：{top['origin']}")
        if "reach_estimate" in top:
            parts.append(f"传播估计：{top['reach_estimate']}")
        if "angles" in top:
            parts.append("角度：" + "; ".join(top.get("angles", [])))
        
        research_text = "\n".join(parts)

        bits = []
        
        # 调用 LLM 生成段子
        from llm.generate_podcast import generate_bit
        bit = generate_bit(meme_name, research_text, prompt.prompt_b1)
        logger.debug(f"Generated bits for meme '{meme_name}': {bit}")
        bits.append(bit)
        bit = generate_bit(meme_name, research_text, prompt.prompt_b2)
        logger.debug(f"Generated bits for meme '{meme_name}': {bit}")
        bits.append(bit)
        bit = generate_bit(meme_name, research_text, prompt.prompt_b3)
        logger.debug(f"Generated bits for meme '{meme_name}': {bit}")
        bits.append(bit)
        content = {"topics": topics, "bits": bits}
        script.content = json.dumps(content, ensure_ascii=False, indent=2)
        db.commit()
        logger.info(f"Saved bits for meme '{meme_name}' with {len(bits)} segments.")
    except Exception as e:
        db.rollback()
        logger.error(f"Failed to generate/save bits: {e}")


def job_generate_script():
    """定时任务：为最新梗生成完整脱口秀脚本并保存至数据库。"""
    logger.debug("Starting job_generate_script")
    db = SessionLocal()
    try:
        # 获取最新的 Script 记录
        script = db.query(Script).filter_by(project="梗文化研究所").order_by(Script.create_time.desc()).first()
        if not script or not script.content:
            logger.warning("No script found for generating full script.")
            return
        
        data = json.loads(script.content)
        topics = data.get("topics", [])
        bits = data.get("bits", [])
        if not topics:
            logger.warning("No topics in the latest script.")
            return
        if not bits:
            logger.warning("No bits in the latest script.")
            return
        
        # 仅处理第一个梗
        top = topics[0]
        meme_name = top.get("title") or top.get("name") or "未知梗"
        logger.debug(f"Generating full script for meme '{meme_name}'")
        # 构建材料文本
        parts = []
        if "summary" in top:
            parts.append(f"简介：{top['summary']}")
        if "origin" in top:
            parts.append(f"可能起源：{top['origin']}")
        if "reach_estimate" in top:
            parts.append(f"传播估计：{top['reach_estimate']}")
        if "angles" in top:
            parts.append("角度：" + "; ".join(top.get("angles", [])))
        
        research_text = "\n".join(parts)
        materials_text = research_text + "\n\n" + json.dumps(bits, ensure_ascii=False, indent=2)
        
        # 调用 LLM 生成完整脚本
        from llm.generate_podcast import generate_script
        full_script = generate_script(meme_name, materials_text)
        content = {"topics": topics, "bits": bits, "script": full_script}
        script.content = json.dumps(content, ensure_ascii=False, indent=2)
        db.commit()
        logger.info(f"Saved full script for meme '{meme_name}'.")
    except Exception as e:
        db.rollback()
        logger.error(f"Failed to generate/save full script: {e}")

# For manual testing
if __name__ == "__main__":
    # job_generate_topics()
    # job_generate_bits()
    job_generate_script()