254 lines
9.6 KiB
Python
254 lines
9.6 KiB
Python
"""
|
||
定时任务调度器
|
||
- 使用 APScheduler 的 AsyncIOScheduler
|
||
- 每小时执行一次检查
|
||
- 查找 queries 表中 status='active' 且 next_query_at <= now() 的记录
|
||
- 为每个符合条件的 query 通过 Agent 框架执行查询
|
||
- 如果 Agent 框架不可用，回退到直接使用 CitationEngine
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||
from apscheduler.triggers.interval import IntervalTrigger
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.database import AsyncSessionLocal
|
||
from app.models.citation_record import CitationRecord
|
||
from app.models.query import Query
|
||
from app.models.query_task import QueryTask
|
||
from app.workers.citation_engine import CitationEngine
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Lazy singleton for CitationDetectorAgent — defers import to avoid
|
||
# circular-dependency issues at module-load time.
|
||
_agent_instance = None
|
||
|
||
|
||
def _get_agent():
|
||
"""Return a cached CitationDetectorAgent instance (lazy singleton)."""
|
||
global _agent_instance
|
||
if _agent_instance is None:
|
||
from app.agent_framework.agents.citation_detector import CitationDetectorAgent
|
||
_agent_instance = CitationDetectorAgent()
|
||
return _agent_instance
|
||
|
||
|
||
class QueryScheduler:
|
||
def __init__(self):
|
||
self.scheduler = AsyncIOScheduler()
|
||
self._agent = None # Lazy-initialized via _get_agent()
|
||
self._fallback_engine = None # Only created when agent fails
|
||
self._loop = None
|
||
|
||
@property
|
||
def agent(self):
|
||
"""Lazy-accessor for CitationDetectorAgent."""
|
||
if self._agent is None:
|
||
self._agent = _get_agent()
|
||
return self._agent
|
||
|
||
@property
|
||
def fallback_engine(self):
|
||
"""Lazy-accessor for fallback CitationEngine (only when agent fails)."""
|
||
if self._fallback_engine is None:
|
||
self._fallback_engine = CitationEngine()
|
||
return self._fallback_engine
|
||
|
||
def start(self):
|
||
"""启动调度器"""
|
||
self._loop = asyncio.get_event_loop()
|
||
self.scheduler.add_job(
|
||
self._run_check,
|
||
trigger=IntervalTrigger(hours=1),
|
||
id="check_queries",
|
||
name="检查并执行到期的查询任务",
|
||
replace_existing=True,
|
||
)
|
||
self.scheduler.add_job(
|
||
self._run_pending_tasks_check,
|
||
trigger=IntervalTrigger(minutes=1),
|
||
id="check_pending_tasks",
|
||
name="检查并执行遗留的pending查询任务",
|
||
replace_existing=True,
|
||
)
|
||
self.scheduler.start()
|
||
logger.info("查询调度器已启动，每小时检查一次待执行任务，每分钟检查一次遗留pending任务")
|
||
|
||
def _run_check(self):
|
||
"""同步包装：将异步检查任务调度到当前事件循环"""
|
||
if self._loop and self._loop.is_running():
|
||
asyncio.run_coroutine_threadsafe(self.check_and_execute_queries(), self._loop)
|
||
else:
|
||
asyncio.run(self.check_and_execute_queries())
|
||
|
||
async def check_and_execute_queries(self):
|
||
"""检查并执行到期的查询任务"""
|
||
logger.info("开始检查待执行查询任务...")
|
||
async with AsyncSessionLocal() as db:
|
||
try:
|
||
now = datetime.now(timezone.utc)
|
||
stmt = select(Query).where(
|
||
Query.status == "active",
|
||
Query.next_query_at <= now,
|
||
)
|
||
result = await db.execute(stmt)
|
||
queries = result.scalars().all()
|
||
|
||
logger.info(f"找到 {len(queries)} 个待执行查询")
|
||
|
||
for query in queries:
|
||
try:
|
||
await self._execute_single_query(query, db)
|
||
except Exception as e:
|
||
logger.error(f"执行查询 {query.id} 失败: {e}")
|
||
continue
|
||
|
||
except Exception as e:
|
||
logger.error(f"检查查询任务时出错: {e}")
|
||
|
||
async def _execute_single_query(self, query: Query, db: AsyncSession):
|
||
"""执行单个查询 — 优先通过 Agent 框架，失败时回退到直接引擎"""
|
||
logger.info(f"开始执行查询: {query.keyword} (ID: {query.id})")
|
||
try:
|
||
await self.agent.execute_query_compat(query, db)
|
||
logger.info(f"查询 {query.id} 执行完成 (via Agent)")
|
||
except Exception as agent_err:
|
||
logger.warning(
|
||
f"Agent 框架执行查询 {query.id} 失败: {agent_err}，回退到直接引擎"
|
||
)
|
||
await self.fallback_engine.execute_query(query, db)
|
||
logger.info(f"查询 {query.id} 执行完成 (via fallback engine)")
|
||
|
||
def _run_pending_tasks_check(self):
|
||
"""同步包装：将异步遗留任务检查调度到当前事件循环"""
|
||
if self._loop and self._loop.is_running():
|
||
asyncio.run_coroutine_threadsafe(self.check_and_execute_pending_tasks(), self._loop)
|
||
else:
|
||
asyncio.run(self.check_and_execute_pending_tasks())
|
||
|
||
async def check_and_execute_pending_tasks(self):
|
||
"""兜底：处理超过1分钟仍未执行的pending任务"""
|
||
logger.info("检查并执行遗留的 pending 查询任务...")
|
||
async with AsyncSessionLocal() as db:
|
||
try:
|
||
one_minute_ago = datetime.utcnow() - timedelta(minutes=1)
|
||
|
||
stmt = select(QueryTask).where(
|
||
QueryTask.status == "pending",
|
||
QueryTask.scheduled_at <= one_minute_ago,
|
||
)
|
||
result = await db.execute(stmt)
|
||
tasks = result.scalars().all()
|
||
|
||
logger.info(f"找到 {len(tasks)} 个遗留的 pending 任务")
|
||
|
||
from collections import defaultdict
|
||
tasks_by_query = defaultdict(list)
|
||
for task in tasks:
|
||
tasks_by_query[task.query_id].append(task)
|
||
|
||
for query_id, task_list in tasks_by_query.items():
|
||
query_stmt = select(Query).where(Query.id == query_id)
|
||
query_result = await db.execute(query_stmt)
|
||
query = query_result.scalar_one_or_none()
|
||
|
||
if not query or query.status != "active":
|
||
continue
|
||
|
||
for task in task_list:
|
||
try:
|
||
task.status = "running"
|
||
task.started_at = datetime.utcnow()
|
||
task.error_message = None
|
||
await db.commit()
|
||
|
||
citation_result = await self._execute_single_platform(
|
||
keyword=query.keyword,
|
||
platform=task.platform,
|
||
target_brand=query.target_brand,
|
||
brand_aliases=query.brand_aliases or [],
|
||
)
|
||
|
||
if citation_result:
|
||
record = CitationRecord.from_citation_result(
|
||
query_id=query_id,
|
||
platform=task.platform,
|
||
result=citation_result,
|
||
)
|
||
db.add(record)
|
||
|
||
task.status = "success"
|
||
task.completed_at = datetime.utcnow()
|
||
await db.commit()
|
||
|
||
except Exception as e:
|
||
await db.rollback()
|
||
task.status = "failed"
|
||
task.error_message = str(e)
|
||
task.completed_at = datetime.utcnow()
|
||
await db.commit()
|
||
logger.error(f"执行遗留任务 {task.id} 失败: {e}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"检查遗留任务时出错: {e}")
|
||
|
||
async def _execute_single_platform(
|
||
self,
|
||
keyword: str,
|
||
platform: str,
|
||
target_brand: str,
|
||
brand_aliases: list,
|
||
) -> dict:
|
||
"""执行单平台检测 — 优先通过 Agent 框架，失败时回退到直接引擎"""
|
||
try:
|
||
result = await self.agent.execute_single_platform_compat(
|
||
keyword=keyword,
|
||
platform=platform,
|
||
target_brand=target_brand,
|
||
brand_aliases=brand_aliases,
|
||
)
|
||
return result
|
||
except Exception as agent_err:
|
||
logger.warning(
|
||
f"Agent 框架执行单平台检测失败 ({platform}): {agent_err}，"
|
||
"回退到直接引擎"
|
||
)
|
||
result = await self.fallback_engine.execute_single_platform(
|
||
keyword=keyword,
|
||
platform=platform,
|
||
target_brand=target_brand,
|
||
brand_aliases=brand_aliases,
|
||
)
|
||
return result
|
||
|
||
async def shutdown(self):
|
||
"""关闭调度器"""
|
||
self.scheduler.shutdown(wait=False)
|
||
if self._agent is not None:
|
||
await self._agent.close()
|
||
if self._fallback_engine is not None:
|
||
await self._fallback_engine.close()
|
||
logger.info("查询调度器已关闭")
|
||
|
||
|
||
# 全局调度器实例
|
||
query_scheduler = QueryScheduler()
|
||
|
||
# 导出别名以兼容测试
|
||
scheduler = query_scheduler.scheduler
|
||
|
||
|
||
def run_job_now(job_id: str):
|
||
"""手动触发指定任务"""
|
||
job = query_scheduler.scheduler.get_job(job_id)
|
||
if job:
|
||
# 获取任务的回调函数并直接调用
|
||
job.func()
|
||
return True
|
||
return False
|