120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
"""World model warmup helpers for startup."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from datetime import datetime, timedelta
|
|
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.exc import OperationalError, ProgrammingError
|
|
|
|
from .core.logging import get_logger
|
|
from .database import AsyncSessionLocal
|
|
from .main_db_guard import is_missing_table_error
|
|
from .models import Project, WorkItem
|
|
from .services.ai_pmo_kernel_service import AIPMOKernelService
|
|
from .utils.timezone import now_cn_aware
|
|
|
|
logger = get_logger(__name__)
|
|
ACTIVE_PROJECT_STATUSES = ("active", "planning", "todo", "doing")
|
|
|
|
|
|
async def get_recent_active_projects(
|
|
*,
|
|
days: int,
|
|
limit_per_tenant: int,
|
|
max_projects: int,
|
|
) -> list[tuple[int, int]]:
|
|
since = (now_cn_aware() - timedelta(days=max(1, int(days)))).replace(tzinfo=None)
|
|
per_tenant_limit = max(1, int(limit_per_tenant))
|
|
total_limit = max(per_tenant_limit, int(max_projects))
|
|
|
|
async with AsyncSessionLocal() as db:
|
|
try:
|
|
rows = (
|
|
await db.execute(
|
|
sa.select(
|
|
WorkItem.tenant_id.label("tenant_id"),
|
|
WorkItem.project_id.label("project_id"),
|
|
sa.func.max(WorkItem.created_at).label("last_seen_at"),
|
|
)
|
|
.join(
|
|
Project,
|
|
sa.and_(
|
|
Project.tenant_id == WorkItem.tenant_id,
|
|
Project.id == WorkItem.project_id,
|
|
),
|
|
)
|
|
.where(
|
|
sa.and_(
|
|
WorkItem.project_id.is_not(None),
|
|
WorkItem.created_at >= since,
|
|
Project.status.in_(ACTIVE_PROJECT_STATUSES),
|
|
)
|
|
)
|
|
.group_by(WorkItem.tenant_id, WorkItem.project_id)
|
|
)
|
|
).all()
|
|
except (OperationalError, ProgrammingError) as exc:
|
|
if not is_missing_table_error(exc):
|
|
raise
|
|
logger.warning("[WM_WARMUP] 活跃项目查询跳过", reason="missing_table", error=str(exc))
|
|
return []
|
|
|
|
bucket: dict[int, list[tuple[datetime, int]]] = {}
|
|
for row in rows:
|
|
tenant_id = int(getattr(row, "tenant_id"))
|
|
project_id = int(getattr(row, "project_id"))
|
|
last_seen = getattr(row, "last_seen_at")
|
|
last_seen_at = last_seen if isinstance(last_seen, datetime) else datetime.min
|
|
bucket.setdefault(tenant_id, []).append((last_seen_at, project_id))
|
|
|
|
picked: list[tuple[datetime, int, int]] = []
|
|
for tenant_id, entries in bucket.items():
|
|
entries.sort(key=lambda item: item[0], reverse=True)
|
|
for last_seen_at, project_id in entries[:per_tenant_limit]:
|
|
picked.append((last_seen_at, tenant_id, project_id))
|
|
|
|
picked.sort(key=lambda item: item[0], reverse=True)
|
|
return [(tenant_id, project_id) for _, tenant_id, project_id in picked[:total_limit]]
|
|
|
|
|
|
async def warmup_world_models() -> None:
|
|
days = max(1, int(os.getenv("WORLD_MODEL_WARMUP_DAYS", "7")))
|
|
limit_per_tenant = max(1, int(os.getenv("WORLD_MODEL_WARMUP_LIMIT_PER_TENANT", "5")))
|
|
max_projects = max(limit_per_tenant, int(os.getenv("WORLD_MODEL_WARMUP_MAX_PROJECTS", "50")))
|
|
try:
|
|
active_projects = await get_recent_active_projects(
|
|
days=days,
|
|
limit_per_tenant=limit_per_tenant,
|
|
max_projects=max_projects,
|
|
)
|
|
warmed = 0
|
|
for tenant_id, project_id in active_projects:
|
|
try:
|
|
async with AsyncSessionLocal() as db:
|
|
kernel = AIPMOKernelService(db, tenant_id=int(tenant_id))
|
|
await kernel.build_world_model(
|
|
tenant_id=int(tenant_id),
|
|
project_id=int(project_id),
|
|
sender="wm_warmup",
|
|
)
|
|
warmed += 1
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"[WM_WARMUP] 项目预热失败",
|
|
tenant_id=tenant_id,
|
|
project_id=project_id,
|
|
error=str(exc),
|
|
)
|
|
logger.info(
|
|
"[WM_WARMUP] 预热完成",
|
|
count=warmed,
|
|
total=len(active_projects),
|
|
days=days,
|
|
limit_per_tenant=limit_per_tenant,
|
|
max_projects=max_projects,
|
|
)
|
|
except Exception as exc:
|
|
logger.warning("[WM_WARMUP] 预热失败", error=str(exc))
|