Agent/examples/content_finder/server.py

"""
内容寻找服务

提供：
1. API 接口：POST /api/tasks - 触发内容寻找任务
2. 定时调度：启动后先恢复 demand_find_task 中 status=执行中 的任务；之后每 2 分钟轮询一次，
   若当前无任务在执行，则从 demand_content 取当天(dt=YYYYMMDD)、未建任务记录且 score 最高的一条执行（不区分品类）
3. 并发控制：限制最大并发任务数；定时侧若已有任务在执行则跳过本次轮询
4. 单次寻找任务最长执行 15 分钟，超时记为失败并回写 demand_find_task
"""

import asyncio
import logging
import os
import uuid
from datetime import datetime
from pathlib import Path
from typing import Optional
import sys

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from zoneinfo import ZoneInfo
from dotenv import load_dotenv

load_dotenv()

import core
from db import (
    create_task_record,
    get_first_running_task,
    get_one_today_unprocessed_demand,
    update_task_status,
    update_task_on_complete,
)
from db.schedule import STATUS_RUNNING, STATUS_SUCCESS, STATUS_FAILED

# 配置日志
log_dir = Path(__file__).parent / '.cache'
log_dir.mkdir(exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(log_dir / 'server.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# FastAPI 应用
app = FastAPI(
    title="内容寻找服务",
    version="1.0.0",
    description="抖音内容寻找 Agent 服务"
)

# 定时调度器（默认用中国时区，避免容器 UTC 导致错过预期时间点）
SCHEDULER_TIMEZONE = os.getenv("SCHEDULER_TIMEZONE", os.getenv("TZ", "Asia/Shanghai"))
scheduler = AsyncIOScheduler(timezone=ZoneInfo(SCHEDULER_TIMEZONE))

# 并发控制
MAX_CONCURRENT_TASKS = int(os.getenv("MAX_CONCURRENT_TASKS", "1"))
task_semaphore = asyncio.Semaphore(MAX_CONCURRENT_TASKS)

# 定时：轮询间隔（分钟）、单次任务超时（秒，默认 15 分钟）
SCHEDULE_INTERVAL_MINUTES = int(os.getenv("SCHEDULE_INTERVAL_MINUTES", "2"))
TASK_TIMEOUT_SECONDS = int(os.getenv("SCHEDULE_TASK_TIMEOUT_SECONDS", "900"))

# 统计信息
stats = {
    "total_tasks": 0,
    "completed_tasks": 0,
    "failed_tasks": 0,
    "scheduled_tasks": 0
}


# ============ 数据模型 ============

class TaskRequest(BaseModel):
    query: Optional[str] = None
    demand_id: Optional[int] = None


class TaskResponse(BaseModel):
    trace_id: str
    status: str
    query: str
    message: str


# ============ 核心函数 ============

def _update_scheduled_task_complete(demand_id: int, trace_id: str, status: int) -> None:
    """定时任务完成时更新 trace_id 和 status，静默处理异常"""
    try:
        update_task_on_complete(demand_id, trace_id, status)
    except Exception as e:
        logger.warning(f"更新任务状态失败: {e}")


async def execute_task(
    query: str,
    demand_id: Optional[int] = None,
    task_type: str = "api",
):
    """
    执行任务（带并发控制）

    Args:
        query: 查询内容
        demand_id: 需求 id（demand_content.id，关联 demand_content 表）
        task_type: 任务类型（"api" 或 "scheduled"）
    """
    async with task_semaphore:
        current_concurrent = MAX_CONCURRENT_TASKS - task_semaphore._value + 1
        logger.info(f"任务开始 [{task_type}]: query={query[:50]}..., 当前并发={current_concurrent}/{MAX_CONCURRENT_TASKS}")

        start_time = datetime.now()
        stats["total_tasks"] += 1
        if task_type == "scheduled":
            stats["scheduled_tasks"] += 1

        if task_type == "scheduled" and demand_id is not None:
            try:
                update_task_status("", demand_id, STATUS_RUNNING)
            except Exception as e:
                logger.warning(f"更新任务状态为执行中失败: {e}")

        try:
            result = await asyncio.wait_for(
                core.run_agent(
                    query, demand_id=demand_id, stream_output=False, log_assistant_text=True
                ),
                timeout=float(TASK_TIMEOUT_SECONDS),
            )
            duration = (datetime.now() - start_time).total_seconds()

            if result["status"] == "completed":
                stats["completed_tasks"] += 1
                logger.info(f"任务完成 [{task_type}]: trace_id={result['trace_id']}, 耗时={duration:.1f}s")
                if task_type == "scheduled" and demand_id is not None:
                    _update_scheduled_task_complete(demand_id, result["trace_id"], STATUS_SUCCESS)
            else:
                stats["failed_tasks"] += 1
                logger.error(f"任务失败 [{task_type}]: trace_id={result.get('trace_id')}, 错误={result.get('error')}, 耗时={duration:.1f}s")
                if task_type == "scheduled" and demand_id is not None:
                    _update_scheduled_task_complete(demand_id, result.get("trace_id") or "", STATUS_FAILED)

        except asyncio.TimeoutError:
            stats["failed_tasks"] += 1
            duration = (datetime.now() - start_time).total_seconds()
            logger.error(
                f"任务超时 [{task_type}]: 超过 {TASK_TIMEOUT_SECONDS}s，记为失败, 耗时={duration:.1f}s"
            )
            if task_type == "scheduled" and demand_id is not None:
                _update_scheduled_task_complete(demand_id, "", STATUS_FAILED)

        except Exception as e:
            stats["failed_tasks"] += 1
            duration = (datetime.now() - start_time).total_seconds()
            logger.error(f"任务异常 [{task_type}]: {e}, 耗时={duration:.1f}s", exc_info=True)
            if task_type == "scheduled" and demand_id is not None:
                _update_scheduled_task_complete(demand_id, "", STATUS_FAILED)


def _today_dt_int() -> int:
    """当天 demand_content.dt 约定为 YYYYMMDD 整数（如 20260402），与定时器时区一致。"""
    return int(datetime.now(ZoneInfo(SCHEDULER_TIMEZONE)).strftime("%Y%m%d"))


def _has_running_content_task() -> bool:
    """
    本进程内是否有内容寻找任务正在执行（占用并发槽）。

    与 execute_task 共用 task_semaphore，含 API 触发与定时触发。
    """
    return task_semaphore._value != MAX_CONCURRENT_TASKS


async def scheduled_tick():
    """
    按 SCHEDULE_INTERVAL_MINUTES 轮询：若当前无任务在执行，则从 demand_content 取
    当天(dt=今日)、尚未出现在 demand_find_task 中且 score 最高的一条需求并执行。
    """
    logger.info("定时任务触发（scheduled_tick）")

    if _has_running_content_task():
        logger.info("定时任务跳过：仍有任务在执行（并发槽占用中）")
        return

    dt = _today_dt_int()
    item = get_one_today_unprocessed_demand(dt=dt)
    if not item:
        logger.info(f"定时任务跳过：无待处理需求（dt={dt} 或均已建任务）")
        return

    demand_content_id = item.get("demand_content_id")
    query = (item.get("query") or "").strip()
    if demand_content_id is None or not query:
        logger.info("定时任务跳过：查询结果无效")
        return

    score = item.get("score")
    logger.info(
        f"定时任务领取（当天 score 最高）：demand_content_id={demand_content_id}, "
        f"dt={dt}, score={score}"
    )
    create_task_record(demand_content_id)
    await execute_task(query=query, demand_id=demand_content_id, task_type="scheduled")


async def run_startup_resume():
    """
    启动后先执行 demand_find_task 中 status=执行中(1) 的任务（理论上仅一条）。
    """
    try:
        row = get_first_running_task()
        if not row:
            logger.info("启动恢复：无执行中(status=1)的 demand_find_task")
            return

        demand_content_id = row.get("demand_content_id")
        query = (row.get("query") or "").strip()
        if demand_content_id is None or not query:
            logger.warning("启动恢复：执行中任务数据不完整，跳过")
            return

        logger.info(f"启动恢复：执行 demand_find_task status=1, demand_content_id={demand_content_id}")
        await execute_task(query=query, demand_id=int(demand_content_id), task_type="scheduled")
    except Exception as e:
        logger.error(f"启动恢复失败: {e}", exc_info=True)


# ============ API 接口 ============

@app.post("/api/tasks", response_model=TaskResponse)
async def create_task(request: TaskRequest):
    """
    创建内容寻找任务

    Args:
        request.query: 查询内容（可选，不传则使用默认值）

    Returns:
        {
            "trace_id": "20260317_103046_xyz789",
            "status": "started",
            "query": "...",
            "message": "任务已启动，结果将保存到 .cache/traces/xxx/"
        }
    """
    # 获取 query 和 demand_id
    query = request.query or core.DEFAULT_QUERY
    demand_id = request.demand_id

    # 用 Event 等待 trace_id
    trace_id_ready = asyncio.Event()
    trace_id_holder = {"id": None}

    async def run_and_capture():
        try:
            # 获取第一个 Trace 对象来获取 trace_id
            from agent import Trace

            async with task_semaphore:
                # 重新构建 runner 来获取 trace_id
                from agent import AgentRunner, RunConfig, FileSystemTraceStore
                from agent.llm import create_openrouter_llm_call
                from agent.llm.prompts import SimplePrompt
                from agent.tools.builtin.knowledge import KnowledgeConfig

                prompt_path = Path(__file__).parent / "content_finder.md"
                prompt = SimplePrompt(prompt_path)
                trace_dir = os.getenv("TRACE_DIR", ".cache/traces")
                demand_id_str = str(demand_id) if demand_id is not None else ""
                messages = prompt.build_messages(query=query, trace_dir=trace_dir, demand_id=demand_id_str)

                api_key = os.getenv("OPEN_ROUTER_API_KEY")
                model_name = prompt.config.get("model", "sonnet-4.6")
                model = os.getenv("MODEL", f"anthropic/claude-{model_name}")
                temperature = float(prompt.config.get("temperature", 0.3))
                max_iterations = int(os.getenv("MAX_ITERATIONS", "30"))
                trace_dir = os.getenv("TRACE_DIR", ".cache/traces")
                skills_dir = str(Path(__file__).parent / "skills")

                Path(trace_dir).mkdir(parents=True, exist_ok=True)
                store = FileSystemTraceStore(base_path=trace_dir)

                allowed_tools = [
                    "douyin_search",
                    "douyin_search_tikhub",
                    "douyin_user_videos",
                    "get_content_fans_portrait",
                    "get_account_fans_portrait",
                    "batch_fetch_portraits",
                    "store_results_mysql",
                    "exec_summary",
                ]

                runner = AgentRunner(
                    llm_call=create_openrouter_llm_call(model=model),
                    trace_store=store,
                    skills_dir=skills_dir,
                )

                config = RunConfig(
                    name="内容寻找",
                    model=model,
                    temperature=temperature,
                    max_iterations=max_iterations,
                    tools=allowed_tools,
                    extra_llm_params={"max_tokens": 8192},
                    knowledge=KnowledgeConfig(
                        enable_extraction=True,
                        enable_completion_extraction=True,
                        enable_injection=True,
                        owner="content_finder_agent",
                        default_tags={"project": "content_finder"},
                        default_scopes=["com.piaoquantv.supply"],
                        default_search_types=["tool", "usecase", "definition"],
                        default_search_owner="content_finder_agent"
                    )
                )

                async for item in runner.run(messages=messages, config=config):
                    if isinstance(item, Trace):
                        if not trace_id_holder["id"]:
                            trace_id_holder["id"] = item.trace_id
                            trace_id_ready.set()
                            logger.info(f"任务启动 [api]: trace_id={item.trace_id}")

                        if item.status == "completed":
                            stats["completed_tasks"] += 1
                            logger.info(f"任务完成 [api]: trace_id={item.trace_id}")
                            break
                        elif item.status == "failed":
                            stats["failed_tasks"] += 1
                            logger.error(f"任务失败 [api]: trace_id={item.trace_id}, 错误={item.error_message}")
                            break

        except Exception as e:
            stats["failed_tasks"] += 1
            logger.error(f"任务异常 [api]: {e}", exc_info=True)
            if not trace_id_holder["id"]:
                trace_id_holder["id"] = f"error_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
                trace_id_ready.set()

    # 启动后台任务
    stats["total_tasks"] += 1
    asyncio.create_task(run_and_capture())

    # 等待 trace_id（最多 5 秒）
    try:
        await asyncio.wait_for(trace_id_ready.wait(), timeout=5.0)
    except asyncio.TimeoutError:
        logger.error("获取 trace_id 超时")
        raise HTTPException(status_code=500, detail="任务启动超时")

    trace_id = trace_id_holder["id"]

    return TaskResponse(
        trace_id=trace_id,
        status="started",
        query=query,
        message=f"任务已启动，结果将保存到 .cache/traces/{trace_id}/"
    )


@app.get("/health")
async def health_check():
    """健康检查"""
    return {
        "status": "ok",
        "max_concurrent_tasks": MAX_CONCURRENT_TASKS,
        "current_tasks": MAX_CONCURRENT_TASKS - task_semaphore._value,
        "scheduler_running": scheduler.running,
        "stats": stats
    }


@app.get("/")
async def root():
    """根路径"""
    return {
        "service": "内容寻找服务",
        "version": "1.0.0",
        "endpoints": {
            "create_task": "POST /api/tasks",
            "health": "GET /health"
        }
    }


# ============ 启动事件 ============

@app.on_event("startup")
async def startup():
    """服务启动时初始化"""
    logger.info("=" * 60)
    logger.info("内容寻找服务启动中...")
    logger.info(f"最大并发任务数: {MAX_CONCURRENT_TASKS}")
    logger.info(f"定时器时区: {SCHEDULER_TIMEZONE}")
    logger.info(
        f"定时策略：每 {SCHEDULE_INTERVAL_MINUTES} 分钟检查是否空闲，空闲则取当天 score 最高的一条；"
        f"单次任务超时 {TASK_TIMEOUT_SECONDS}s"
    )

    asyncio.create_task(run_startup_resume())

    job = scheduler.add_job(
        scheduled_tick,
        "interval",
        minutes=SCHEDULE_INTERVAL_MINUTES,
        misfire_grace_time=300,
        coalesce=True,
        max_instances=1,
    )
    scheduler.start()
    logger.info(f"定时任务已注册: id={job.id}, next_run_time={job.next_run_time}")

    logger.info("服务启动完成")
    logger.info("=" * 60)


@app.on_event("shutdown")
async def shutdown():
    """服务关闭时清理"""
    logger.info("服务关闭中...")
    if scheduler.running:
        scheduler.shutdown()
    logger.info("服务已关闭")


# ============ 主函数 ============

if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("PORT", "8080"))
    host = os.getenv("HOST", "0.0.0.0")

    logger.info(f"启动服务: http://{host}:{port}")
    uvicorn.run(app, host=host, port=port)