"""定时任务执行器。 一次执行 = 建一个独立会话 + 跑一轮 agent + 把结果落库并计费。 为什么用「独立会话」而不是 1.5 那种 ``scheduled-task-{id}-{run_id}`` 裸线程 ID: 本系统其他所有地方都强制 ``thread_id ≡ session_id``(那条不变式是防止 跨会话读写图状态的第一道闸)。定时任务若用另一个 ID 命名法,就成了唯一 的例外,将来谁改了都会踩。用独立会话既同样隔离,又让用户能在会话列表里 直接翻到定时任务的产出——1.5 的定时任务结果只能去 schedule_runs 表里看。 """ from __future__ import annotations import asyncio from typing import Any from langchain_core.messages import AIMessage, HumanMessage import config from api.context_meter import extract_usage from billing import pricing from billing.meter import UsageMeter from db import chat_store, db_conn, schedule_store from obs import get_logger logger = get_logger(__name__) # 定时任务的图递归上限:比交互式对话紧得多。 # 理由是没人盯着它——跑飞了只会白烧额度,早点中止更划算。 SCHEDULED_RECURSION_LIMIT = 30 def _extract_final_text(result: dict[str, Any]) -> str: """从图输出里取最后一条 AI 回复的正文。""" for message in reversed(result.get("messages") or []): if isinstance(message, AIMessage): content = message.content if isinstance(content, str) and content.strip(): return content if isinstance(content, list): text = "".join( block.get("text", "") for block in content if isinstance(block, dict) ) if text.strip(): return text return "" def _sum_usage(messages: list[Any]) -> tuple[int, int, str]: """汇总本次执行的真实 token 用量。""" input_tokens = output_tokens = 0 source = "provider" saw_provider = False for message in messages or []: usage = extract_usage(message) if usage is None: continue saw_provider = True input_tokens += usage["input_tokens"] output_tokens += usage["output_tokens"] if not saw_provider and messages: # provider 没给用量:标 estimate 而不是静默按 0 记账 source = "estimate" return input_tokens, output_tokens, source async def run_schedule(schedule_id: str, trigger: str = "cron") -> dict[str, Any]: """执行一次定时任务。异常不外抛——调度器的回调抛异常会被 APScheduler 吞掉并只留一行日志,不如自己接住并写进 schedule_runs 让用户看得见。""" schedule = await asyncio.to_thread(schedule_store.get_schedule, schedule_id) if schedule is None: logger.warning("定时任务不存在 id=%s", schedule_id) return {"ok": False, "error": "schedule_not_found"} user_name = schedule.get("user_name") or "default" title = f"[定时] {schedule.get('title') or schedule_id[:8]}" session = await asyncio.to_thread( chat_store.create_session, user_name, schedule.get("mode") or "full", title ) thread_id = session["id"] run_id = await asyncio.to_thread(schedule_store.start_run, schedule_id, thread_id) logger.info( "定时任务开始 id=%s run=%s trigger=%s thread=%s", schedule_id, run_id[:8], trigger, thread_id, ) prompt = schedule.get("prompt") or "" await asyncio.to_thread(chat_store.save_message, thread_id, "user", prompt) try: from agents import agent_registry agent = agent_registry.get_agent() result = await agent.ainvoke( {"messages": [HumanMessage(content=prompt)]}, config={ "configurable": {"thread_id": thread_id}, "recursion_limit": SCHEDULED_RECURSION_LIMIT, }, ) except Exception as exc: detail = f"{type(exc).__name__}: {exc}" logger.exception("定时任务失败 id=%s run=%s", schedule_id, run_id[:8]) await asyncio.to_thread( chat_store.save_message, thread_id, "assistant", f"(任务执行失败){detail}" ) await asyncio.to_thread(schedule_store.finish_run, run_id, "failed", detail, None) return {"ok": False, "run_id": run_id, "session_id": thread_id, "error": detail} output = _extract_final_text(result) await asyncio.to_thread(chat_store.save_message, thread_id, "assistant", output) await asyncio.to_thread(schedule_store.finish_run, run_id, "success", None, output[:2000]) # 计费:定时任务同样消耗额度,不记就是白跑 input_tokens, output_tokens, source = _sum_usage(result.get("messages") or []) if input_tokens or output_tokens: meter = UsageMeter( user_name=user_name, session_id=thread_id, kind="scheduled", model=config.MODEL, ) meter.add(input_tokens, output_tokens, source) meter.settle() logger.info( "定时任务完成 id=%s run=%s 输出长度=%d tokens=%d/%d", schedule_id, run_id[:8], len(output), input_tokens, output_tokens, ) return { "ok": True, "run_id": run_id, "session_id": thread_id, "output": output, "input_tokens": input_tokens, "output_tokens": output_tokens, } def ensure_pricing_seeded() -> None: """确保单价表就绪,避免定时任务因为模型没定价而 fail-closed。""" try: pricing.seed_default_pricing() except Exception as exc: # pragma: no cover - 取决于数据库可用性 logger.warning("单价表初始化失败:%s: %s", type(exc).__name__, exc) __all__ = ["SCHEDULED_RECURSION_LIMIT", "ensure_pricing_seeded", "run_schedule", "db_conn"]