浅念
259 lines
10 KiB
Python
259 lines
10 KiB
Python
"""site_import 内存任务表（spec §3.3）——单进程 asyncio，(slug, task_id) 索引。
|
||
|
||
内存铁律：30 分钟惰性清理（查询/新建时顺带扫）；强引用防 GC（裸
|
||
create_task 会被回收——asyncio 官方坑，模块级 dict 存 task 对象）；
|
||
全平台并发 ≤2（复活任务占用同一名额）；频控 3 次/15 分钟（新建与
|
||
retry 同计）。异常类型：ImportRateLimited→429 / ConcurrentLimit→409 /
|
||
LookupError→404 / ValueError→409——API 层（Task 7）映射。
|
||
|
||
7 分钟总超时兜底：超时落在生成段（抓取产物已在）标 failed_phase=
|
||
"generate"——retry=generate 是 spec §3.3 的预期恢复路径（不重抓）；
|
||
落在抓取段则无产物，重开任务。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import time
|
||
import uuid
|
||
from collections import defaultdict
|
||
|
||
from app.platform.site_import import crawler, generator
|
||
from app.platform.site_import.crawler import CrawlError, INSUFFICIENT_MESSAGE
|
||
from app.platform.site_import.generator import GenerateError
|
||
from app.platform.site_import.safe_fetch import FetchError
|
||
|
||
TASK_TTL = 30 * 60
|
||
TOTAL_TIMEOUT = 7 * 60
|
||
MAX_CONCURRENT = 2
|
||
IMPORT_WINDOW = 15 * 60
|
||
IMPORT_MAX = 3
|
||
|
||
|
||
class ImportRateLimited(Exception):
|
||
"""频控超限（15 分钟 3 次）——API 层映射 429。"""
|
||
|
||
|
||
class ConcurrentLimit(Exception):
|
||
"""全平台并发上限（≤2，复活任务同占名额）——API 层映射 409。"""
|
||
|
||
|
||
class _SlidingWindow:
|
||
"""(kind, slug) 键滑动窗计数——config_routes.SlidingWindowCounter 同款
|
||
风格的独立实例（勿共用 email/push 测试端点的预算）。同款语义：超限后
|
||
allow 守卫先于一切预检；计数只落在真正发起的这次（校验失败不消耗）。"""
|
||
|
||
def __init__(self, window_seconds: float = IMPORT_WINDOW,
|
||
max_calls: int = IMPORT_MAX) -> None:
|
||
self.window_seconds = window_seconds
|
||
self.max_calls = max_calls
|
||
self._calls: dict[tuple[str, str], list[float]] = defaultdict(list)
|
||
|
||
def allow(self, key: tuple[str, str]) -> bool:
|
||
self._prune(key, time.time())
|
||
return len(self._calls.get(key, ())) < self.max_calls
|
||
|
||
def record(self, key: tuple[str, str]) -> None:
|
||
now = time.time()
|
||
self._prune(key, now)
|
||
self._calls[key].append(now)
|
||
|
||
def _prune(self, key: tuple[str, str], now: float) -> None:
|
||
calls = self._calls.get(key)
|
||
if calls is None:
|
||
return
|
||
kept = [t for t in calls if t > now - self.window_seconds]
|
||
if kept:
|
||
self._calls[key] = kept
|
||
else:
|
||
self._calls.pop(key, None)
|
||
|
||
|
||
# 模块级单例：导入频控（键 ("import", slug)，全平台各租户独立预算）
|
||
_import_limiter = _SlidingWindow()
|
||
|
||
# 内存任务表：(slug, task_id) 联合索引——跨租户持他人 id 查询即 miss
|
||
_tasks: dict[tuple[str, str], "ImportTask"] = {}
|
||
_lock = asyncio.Lock()
|
||
|
||
|
||
class TaskState:
|
||
PENDING = "pending"
|
||
CRAWLING = "crawling"
|
||
GENERATING = "generating"
|
||
DONE = "done"
|
||
FAILED = "failed"
|
||
|
||
ACTIVE = (PENDING, CRAWLING, GENERATING)
|
||
|
||
|
||
class ImportTask:
|
||
"""一条导入任务：状态机 + 阶段产物 + 强引用的后台 aio task。
|
||
|
||
failed_phase："generate" 可 retry=generate 复活（复用 crawl_result，
|
||
不重抓）；"crawl"（SPA/内容不足/FetchError 网站打不开）与 None（超时
|
||
落在抓取段/未知异常）不可复活——重开任务。
|
||
"""
|
||
|
||
def __init__(self, slug: str, task_id: str, url: str | None,
|
||
config_data: dict | None = None) -> None:
|
||
self.slug = slug
|
||
self.task_id = task_id
|
||
self.url = url
|
||
# 租户配置快照（终审 Important #2：AI resolver 走三级链——
|
||
# 租户自有 ai.* 配置优先于平台键）
|
||
self.tenant_config_data = config_data
|
||
self.state = TaskState.PENDING
|
||
self.phase: str | None = None # 运行中阶段（crawling/generating）
|
||
self.pages_fetched = 0
|
||
self.draft: dict | None = None
|
||
self.error: str | None = None
|
||
self.created_at = time.time()
|
||
self.finished_at: float | None = None
|
||
self.failed_phase: str | None = None
|
||
self.crawl_result: dict | None = None # retry=generate 复用
|
||
self.need_crawl = True
|
||
self._aio_task: asyncio.Task | None = None # 强引用防 GC
|
||
|
||
def fail(self, message: str, *, phase: str | None) -> None:
|
||
self.state = TaskState.FAILED
|
||
self.error = message
|
||
self.failed_phase = phase
|
||
self.phase = None
|
||
self.finished_at = time.time()
|
||
|
||
def _finish(self, draft: dict) -> None:
|
||
self.draft = draft
|
||
self.state = TaskState.DONE
|
||
self.phase = None
|
||
self.finished_at = time.time()
|
||
|
||
|
||
def _prune_expired(now: float | None = None) -> None:
|
||
"""30 分钟惰性清理（查询/新建顺带扫）。运行中任务防御性取消——
|
||
7 分钟总超时下正常不会发生（30 >> 7），仅兜异常态。"""
|
||
now = time.time() if now is None else now
|
||
deadline = now - TASK_TTL
|
||
for key in [k for k, t in _tasks.items() if t.created_at < deadline]:
|
||
stale = _tasks.pop(key)
|
||
if stale._aio_task is not None and not stale._aio_task.done():
|
||
stale._aio_task.cancel()
|
||
|
||
|
||
def _count_active() -> int:
|
||
return sum(1 for t in _tasks.values() if t.state in TaskState.ACTIVE)
|
||
|
||
|
||
def active_count() -> int:
|
||
"""全平台进行中任务数（pending/crawling/generating）。"""
|
||
_prune_expired()
|
||
return _count_active()
|
||
|
||
|
||
def get(slug: str, task_id: str) -> ImportTask | None:
|
||
"""(slug, task_id) 联合索引查询 + 惰性清理（跨租户 → None → 404）。"""
|
||
_prune_expired()
|
||
return _tasks.get((slug, task_id))
|
||
|
||
|
||
async def start(slug: str, url: str | None, *,
|
||
retry_task_id: str | None = None,
|
||
config_data: dict | None = None) -> str:
|
||
"""新建或复活（retry=generate）导入任务，返回 task_id。
|
||
|
||
频控守卫最先（超限后 429 优先于一切预检）；record 只落在真正发起的
|
||
这次——LookupError/ValueError/ConcurrentLimit 不消耗预算。
|
||
config_data：租户配置快照（AI resolver 三级链的租户槽位；None 走
|
||
平台设置→env——新建空租户即此形态）。
|
||
"""
|
||
_prune_expired()
|
||
key = ("import", slug)
|
||
if not _import_limiter.allow(key):
|
||
raise ImportRateLimited(
|
||
f"导入太频繁，{IMPORT_WINDOW // 60} 分钟内最多 {IMPORT_MAX} 次，"
|
||
"稍后再试")
|
||
async with _lock:
|
||
if retry_task_id is not None:
|
||
target = _tasks.get((slug, retry_task_id))
|
||
if target is None:
|
||
raise LookupError(
|
||
f"task {retry_task_id!r} not found for tenant {slug!r}")
|
||
if (target.state != TaskState.FAILED
|
||
or target.failed_phase != "generate"):
|
||
raise ValueError("任务当前状态不可重试"
|
||
"（仅 generate 阶段失败可复活）")
|
||
# 并发名额对新建与 retry 复活同样生效（spec §3.3）
|
||
if _count_active() >= MAX_CONCURRENT:
|
||
raise ConcurrentLimit("已有导入任务进行中，稍后再试")
|
||
if retry_task_id is not None:
|
||
# 复活同 task_id：复用已抓产物，不重抓
|
||
target.state = TaskState.PENDING
|
||
target.error = None
|
||
target.failed_phase = None
|
||
target.finished_at = None
|
||
target.need_crawl = False
|
||
else:
|
||
target = ImportTask(slug, uuid.uuid4().hex, url,
|
||
config_data=config_data)
|
||
_tasks[(slug, target.task_id)] = target
|
||
_import_limiter.record(key)
|
||
target._aio_task = asyncio.create_task(_pipeline(target))
|
||
return target.task_id
|
||
|
||
|
||
async def _pipeline(task_obj: ImportTask) -> None:
|
||
"""后台管线：7 分钟总超时兜底 + 异常归档为白话失败。"""
|
||
try:
|
||
await asyncio.wait_for(_run(task_obj), TOTAL_TIMEOUT)
|
||
except asyncio.TimeoutError:
|
||
task_obj.fail(
|
||
"导入超时，请重试",
|
||
phase="generate" if task_obj.crawl_result is not None else None)
|
||
except CrawlError as exc:
|
||
task_obj.fail(str(exc), phase="crawl") # SPA/内容不足——不可 retry
|
||
except FetchError:
|
||
# spec §6 行1：入口页不可达/超时是最高频真实失败路径——白话文案，
|
||
# 不漏进通用兜底裸抛英文内部信息
|
||
task_obj.fail("网站打不开或响应太慢——检查 URL", phase="crawl")
|
||
except GenerateError as exc:
|
||
task_obj.fail(str(exc), phase="generate") # retry=generate 复活
|
||
except Exception as exc: # 防御：未知异常不裸抛（任务表对象完整性优先）
|
||
task_obj.fail(f"导入失败：{exc}", phase=None)
|
||
|
||
|
||
async def _run(task_obj: ImportTask) -> None:
|
||
if task_obj.need_crawl:
|
||
task_obj.state = TaskState.CRAWLING
|
||
task_obj.phase = "crawling"
|
||
task_obj.crawl_result = await crawler.crawl(task_obj.url)
|
||
task_obj.pages_fetched = int(
|
||
task_obj.crawl_result.get("pages_fetched") or 0)
|
||
task_obj.state = TaskState.GENERATING
|
||
task_obj.phase = "generating"
|
||
# resolver 三级链：租户 ai 组 → 平台设置 → env（终审 Important #2：
|
||
# 传租户配置快照——自有 AI 配置的租户导入烧自己的配额，不用平台键）
|
||
provider = generator.make_provider(task_obj.tenant_config_data)
|
||
# source_text 传 analyze 实际消费的截断后文本（24K 同款——Task 5 备忘：
|
||
# 否则截断段数字会被 build 的数字校验误剔除）
|
||
source = ((task_obj.crawl_result or {}).get("pages_text")
|
||
or "")[:generator.TOTAL_INPUT_CAP]
|
||
analysis = await generator.analyze(provider, source)
|
||
if not (analysis.get("product_lines") or []):
|
||
# 内容不足三触发之一（调用 1 产品线为空）——同 crawl 类不可复活
|
||
task_obj.fail(INSUFFICIENT_MESSAGE, phase="crawl")
|
||
return
|
||
gen = await generator.generate(provider, analysis)
|
||
task_obj._finish(
|
||
generator.build(gen, analysis, task_obj.crawl_result, source))
|
||
|
||
|
||
def _reset_for_test() -> None:
|
||
"""测试隔离（conftest autouse 前后各跑）：任务表清空（运行中 aio
|
||
任务防御性取消）+ 频控窗计数/参数复位。"""
|
||
for t in _tasks.values():
|
||
if t._aio_task is not None and not t._aio_task.done():
|
||
t._aio_task.cancel()
|
||
_tasks.clear()
|
||
_import_limiter._calls.clear()
|
||
_import_limiter.window_seconds = IMPORT_WINDOW
|
||
_import_limiter.max_calls = IMPORT_MAX
|