485 lines
18 KiB
Python
485 lines
18 KiB
Python
"""PMFlow FastAPI 应用入口。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import os
|
||
import time
|
||
from contextlib import asynccontextmanager, suppress
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
from fastapi import FastAPI
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
import sqlalchemy as sa
|
||
from sqlalchemy.exc import OperationalError, ProgrammingError
|
||
|
||
from .api.router import api_router
|
||
from .core.config import settings
|
||
from .core.logging import configure_logging, get_logger
|
||
from .core.response import success_response
|
||
from .database import AsyncSessionLocal, engine
|
||
from .middleware import AuditContextMiddleware
|
||
from .models import Project, WorkItem
|
||
from .services.scheduler_center import WorkloadRecalcScheduler
|
||
from .services.scheduler_center import KnowledgeUrlRefreshScheduler
|
||
from .services.scheduler_center import KnowledgeGraphIncrementalScheduler
|
||
from .services.scheduler_center import AIPMOCadenceScheduler
|
||
from .services.ai_orchestration_run_service import AIOrchestrationRunService
|
||
from .services.ai_pmo_kernel_service import AIPMOKernelService
|
||
from .services.skill_service import init_system_skill_tags
|
||
from .services.feishu_long_connection import init_feishu_long_connection_manager
|
||
from .services.feishu_integration import dispatch_pending_feishu_replies
|
||
from .services.context_cache import register_context_cache_invalidation_hooks
|
||
|
||
from .utils.timezone import now_cn_aware
|
||
configure_logging(settings.debug)
|
||
logger = get_logger(__name__)
|
||
register_context_cache_invalidation_hooks()
|
||
|
||
ACTIVE_PROJECT_STATUSES = ("active", "planning", "todo", "doing")
|
||
|
||
|
||
def _is_missing_table_error(exc: Exception) -> bool:
|
||
text = str(exc).lower()
|
||
return (
|
||
"no such table" in text
|
||
or "undefined table" in text
|
||
or "relation" in text and "does not exist" in text
|
||
)
|
||
|
||
|
||
def _normalize_dialect(value: str) -> str:
|
||
raw = str(value or "").strip().lower()
|
||
if not raw:
|
||
return ""
|
||
if raw.startswith("postgres"):
|
||
return "postgresql"
|
||
if raw.startswith("mysql"):
|
||
return "mysql"
|
||
if "+" in raw:
|
||
raw = raw.split("+", 1)[0]
|
||
if ":" in raw:
|
||
raw = raw.split(":", 1)[0]
|
||
return raw
|
||
|
||
|
||
def _mask_db_url(url: str) -> str:
|
||
value = str(url or "").strip()
|
||
if "@" not in value:
|
||
return value
|
||
prefix, suffix = value.split("@", 1)
|
||
if "://" not in prefix:
|
||
return f"***@{suffix}"
|
||
scheme, creds = prefix.split("://", 1)
|
||
user = creds.split(":", 1)[0] if ":" in creds else creds
|
||
return f"{scheme}://{user}:***@{suffix}"
|
||
|
||
|
||
def _check_database_dialect_guard(sync_conn) -> None:
|
||
configured = _normalize_dialect(settings.database_url)
|
||
actual = _normalize_dialect(sync_conn.dialect.name)
|
||
expected = _normalize_dialect(os.getenv("PMFLOW_EXPECTED_DB_DIALECT", "postgresql")) or "postgresql"
|
||
|
||
env_urls = [
|
||
os.getenv("DATABASE_URL", ""),
|
||
os.getenv("SQLALCHEMY_DATABASE_URL", ""),
|
||
os.getenv("SYNC_DATABASE_URL", ""),
|
||
]
|
||
configured_families = {_normalize_dialect(item) for item in env_urls if str(item or "").strip()}
|
||
configured_families.discard("")
|
||
|
||
logger.info(
|
||
"database_dialect_guard",
|
||
configured=configured,
|
||
actual=actual,
|
||
expected=expected or "unset",
|
||
database_url=_mask_db_url(settings.database_url),
|
||
)
|
||
print(
|
||
"[DB-Guard] "
|
||
f"configured={configured}, actual={actual}, expected={expected or 'unset'}, "
|
||
f"url={_mask_db_url(settings.database_url)}"
|
||
)
|
||
if configured != "postgresql":
|
||
raise RuntimeError(f"DATABASE_URL 必须为 PostgreSQL，当前 configured={configured or 'unset'}。")
|
||
if actual != "postgresql":
|
||
raise RuntimeError(f"运行时数据库必须为 PostgreSQL，当前 actual={actual or 'unset'}。")
|
||
if configured and actual and configured != actual:
|
||
raise RuntimeError(
|
||
f"数据库方言不一致: configured={configured}, actual={actual}。"
|
||
"请确保同一服务只连接单一数据库方言。"
|
||
)
|
||
if expected and actual and expected != actual:
|
||
raise RuntimeError(
|
||
f"数据库方言与 PMFLOW_EXPECTED_DB_DIALECT 不一致: expected={expected}, actual={actual}。"
|
||
)
|
||
if len(configured_families) > 1:
|
||
raise RuntimeError(
|
||
f"检测到多种数据库方言配置: {sorted(configured_families)}。"
|
||
"禁止同一服务混用多数据库方言。"
|
||
)
|
||
|
||
|
||
def _is_postgresql_runtime() -> bool:
|
||
return _normalize_dialect(engine.dialect.name) == "postgresql"
|
||
|
||
|
||
async def _feishu_outbox_polling_loop() -> None:
|
||
while True:
|
||
try:
|
||
dispatched = await dispatch_pending_feishu_replies(limit=50)
|
||
await asyncio.sleep(0.3 if dispatched > 0 else 2.0)
|
||
except asyncio.CancelledError:
|
||
break
|
||
except Exception as exc:
|
||
logger.warning("feishu_outbox_drain_failed", error=str(exc))
|
||
await asyncio.sleep(2.0)
|
||
|
||
|
||
async def _feishu_outbox_listen_loop() -> None:
|
||
while True:
|
||
listener_callback = None
|
||
driver_connection = None
|
||
try:
|
||
async with engine.connect() as conn:
|
||
await conn.execute(sa.text("LISTEN feishu_outbox_ready"))
|
||
await conn.commit()
|
||
raw_conn = await conn.get_raw_connection()
|
||
driver_connection = getattr(raw_conn, "driver_connection", None)
|
||
if driver_connection is None or not hasattr(driver_connection, "add_listener"):
|
||
logger.warning("feishu_outbox_listen_unsupported_driver", driver=str(type(driver_connection)))
|
||
await _feishu_outbox_polling_loop()
|
||
return
|
||
|
||
notify_queue: asyncio.Queue[str] = asyncio.Queue(maxsize=1024)
|
||
|
||
def _on_notify(_connection, _pid, _channel, payload) -> None:
|
||
try:
|
||
notify_queue.put_nowait(str(payload or ""))
|
||
except asyncio.QueueFull:
|
||
pass
|
||
|
||
listener_callback = _on_notify
|
||
await driver_connection.add_listener("feishu_outbox_ready", listener_callback)
|
||
logger.info("feishu_outbox_listen_ready")
|
||
|
||
while True:
|
||
dispatched = await dispatch_pending_feishu_replies(limit=50)
|
||
if dispatched > 0:
|
||
continue
|
||
try:
|
||
await asyncio.wait_for(notify_queue.get(), timeout=2.0)
|
||
while not notify_queue.empty():
|
||
notify_queue.get_nowait()
|
||
except asyncio.TimeoutError:
|
||
continue
|
||
except asyncio.CancelledError:
|
||
break
|
||
except Exception as exc:
|
||
logger.warning("feishu_outbox_listen_failed", error=str(exc))
|
||
await asyncio.sleep(2.0)
|
||
finally:
|
||
if driver_connection is not None and listener_callback is not None:
|
||
with suppress(Exception):
|
||
await driver_connection.remove_listener("feishu_outbox_ready", listener_callback)
|
||
|
||
|
||
async def _feishu_outbox_drain_loop() -> None:
|
||
if _is_postgresql_runtime():
|
||
await _feishu_outbox_listen_loop()
|
||
return
|
||
await _feishu_outbox_polling_loop()
|
||
|
||
|
||
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))
|
||
|
||
|
||
def _startup_log(message: str, **fields) -> None:
|
||
logger.info(f"[后端] {message}", **fields)
|
||
|
||
|
||
async def _probe_model_connectivity_on_startup() -> None:
|
||
from .services.llm import llm_service
|
||
|
||
low_provider = str(os.getenv("LLM_ROUTING_LOW_PROVIDER", os.getenv("LLM_PROVIDER", "glm")) or "").strip() or "glm"
|
||
low_model = str(os.getenv("LLM_ROUTING_LOW_MODEL", os.getenv("GLM_MODEL", "GLM-4.7")) or "").strip() or "GLM-4.7"
|
||
|
||
high_provider = str(os.getenv("LLM_ROUTING_HIGH_PROVIDER", os.getenv("LLM_PROVIDER", "openai")) or "").strip() or "openai"
|
||
high_model = str(os.getenv("LLM_ROUTING_HIGH_MODEL", os.getenv("OPENAI_MODEL", "GLM-5")) or "").strip() or "GLM-5"
|
||
|
||
probe_targets: list[tuple[str, str, str]] = [
|
||
("LOW", low_provider, low_model),
|
||
("HIGH", high_provider, high_model),
|
||
]
|
||
deduped_targets: list[tuple[str, str, str]] = []
|
||
seen: set[tuple[str, str]] = set()
|
||
for level, provider, model in probe_targets:
|
||
key = (provider.lower(), model)
|
||
if key in seen:
|
||
continue
|
||
seen.add(key)
|
||
deduped_targets.append((level, provider, model))
|
||
|
||
lines: list[str] = [
|
||
"┌──────────────────────────────┐",
|
||
"│ 模型连通性检测 │",
|
||
]
|
||
|
||
for level, provider, model in deduped_targets:
|
||
try:
|
||
runtime = llm_service._env_provider_runtime(provider)
|
||
runtime["model"] = model
|
||
except Exception as exc:
|
||
logger.error(
|
||
"startup_model_probe_runtime_invalid",
|
||
level=level,
|
||
provider=provider,
|
||
model=model,
|
||
error=str(exc),
|
||
)
|
||
lines.append(f"│ {model:<12} ✗ 失败 (0.0s) │")
|
||
continue
|
||
started_at = time.perf_counter()
|
||
ok = False
|
||
error_text = ""
|
||
try:
|
||
result = await asyncio.wait_for(
|
||
llm_service._chat_completion(
|
||
provider=str(runtime.get("provider") or provider),
|
||
prompt="ping",
|
||
system_prompt="你是连通性检测助手，只回复 OK。",
|
||
api_key=str(runtime.get("api_key") or ""),
|
||
base_url=str(runtime.get("base_url") or ""),
|
||
model=str(runtime.get("model") or model),
|
||
temperature=0.0,
|
||
max_tokens=1,
|
||
function_name="startup_model_probe",
|
||
response_mode="fast",
|
||
),
|
||
timeout=8.0,
|
||
)
|
||
ok = bool(str(result or "").strip())
|
||
if not ok:
|
||
error_text = "empty_response"
|
||
except Exception as exc:
|
||
error_text = str(exc)
|
||
logger.error(
|
||
"startup_model_probe_failed",
|
||
level=level,
|
||
provider=provider,
|
||
model=model,
|
||
error=error_text,
|
||
)
|
||
|
||
elapsed = time.perf_counter() - started_at
|
||
if ok:
|
||
lines.append(f"│ {model:<12} ✓ 正常 ({elapsed:.1f}s) │")
|
||
else:
|
||
lines.append(f"│ {model:<12} ✗ 失败 ({elapsed:.1f}s) │")
|
||
if error_text:
|
||
_startup_log("模型连通性失败（不阻断启动）", model=model, provider=provider, error=error_text)
|
||
|
||
lines.append("└──────────────────────────────┘")
|
||
for line in lines:
|
||
logger.info(line)
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(_: FastAPI):
|
||
"""应用生命周期：初始化表与系统技能标签。"""
|
||
backend_port = str(os.getenv("BACKEND_PORT", "8000") or "8000").strip() or "8000"
|
||
_startup_log("服务启动初始化中")
|
||
auto_create_tables = os.getenv("AUTO_CREATE_TABLES", "false").lower() in {
|
||
"1",
|
||
"true",
|
||
"yes",
|
||
}
|
||
if auto_create_tables:
|
||
logger.error(
|
||
"runtime_ddl_blocked",
|
||
auto_create_tables=auto_create_tables,
|
||
reason="runtime ddl is forbidden, run alembic migration instead",
|
||
)
|
||
raise RuntimeError("AUTO_CREATE_TABLES=true 已被禁用；请使用 Alembic migration 变更数据库。")
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(_check_database_dialect_guard)
|
||
_startup_log("数据库连接成功")
|
||
await _probe_model_connectivity_on_startup()
|
||
|
||
async with AsyncSessionLocal() as db:
|
||
try:
|
||
created = await init_system_skill_tags(db, tenant_id=1)
|
||
if created:
|
||
logger.info("system_skill_tags_seeded", created=created)
|
||
except Exception as exc: # pragma: no cover
|
||
logger.warning("seed_system_skills_failed", error=str(exc))
|
||
try:
|
||
interrupted = await AIOrchestrationRunService(db).mark_stale_runs_interrupted(older_than_minutes=10)
|
||
if interrupted:
|
||
await db.commit()
|
||
logger.info("stale_ai_orchestration_runs_interrupted", count=interrupted)
|
||
except Exception as exc: # pragma: no cover
|
||
logger.warning("interrupt_stale_ai_orchestration_runs_failed", error=str(exc))
|
||
|
||
scheduler = WorkloadRecalcScheduler(AsyncSessionLocal)
|
||
scheduler.start()
|
||
knowledge_scheduler = KnowledgeUrlRefreshScheduler(AsyncSessionLocal)
|
||
knowledge_scheduler.start()
|
||
kg_scheduler = KnowledgeGraphIncrementalScheduler(AsyncSessionLocal)
|
||
kg_scheduler.start()
|
||
ai_pmo_cadence_scheduler = AIPMOCadenceScheduler(AsyncSessionLocal)
|
||
ai_pmo_cadence_scheduler.start()
|
||
feishu_lc = init_feishu_long_connection_manager(AsyncSessionLocal)
|
||
await feishu_lc.start()
|
||
outbox_task = asyncio.create_task(
|
||
_feishu_outbox_drain_loop(),
|
||
name="feishu-outbox-drain-loop",
|
||
)
|
||
wm_warmup_task = asyncio.create_task(
|
||
_warmup_world_models(),
|
||
name="wm-world-model-warmup",
|
||
)
|
||
_startup_log("后端运行中", url=f"http://127.0.0.1:{backend_port}")
|
||
|
||
yield
|
||
_startup_log("服务关闭中")
|
||
outbox_task.cancel()
|
||
try:
|
||
await outbox_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
if not wm_warmup_task.done():
|
||
wm_warmup_task.cancel()
|
||
with suppress(asyncio.CancelledError):
|
||
await wm_warmup_task
|
||
await feishu_lc.stop()
|
||
await ai_pmo_cadence_scheduler.stop()
|
||
await kg_scheduler.stop()
|
||
await knowledge_scheduler.stop()
|
||
await scheduler.stop()
|
||
await engine.dispose()
|
||
_startup_log("服务已关闭")
|
||
|
||
|
||
app = FastAPI(
|
||
title=settings.app_name,
|
||
debug=settings.debug,
|
||
version="0.1.0",
|
||
lifespan=lifespan,
|
||
)
|
||
|
||
app.add_middleware(AuditContextMiddleware)
|
||
|
||
origins = [item.strip() for item in settings.cors_origins.split(",") if item.strip()]
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=origins or ["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
app.include_router(api_router)
|
||
|
||
|
||
@app.get("/healthz")
|
||
async def healthz() -> dict:
|
||
return success_response({"status": "ok"})
|