213 lines
7.1 KiB
Python
213 lines
7.1 KiB
Python
import uuid
|
||
import logging
|
||
from datetime import datetime, timezone
|
||
from celery import Celery
|
||
from celery.schedules import crontab
|
||
|
||
from .core.config import settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Celery 配置
|
||
celery_app = Celery(
|
||
"compliance",
|
||
broker=settings.redis_url,
|
||
backend=settings.redis_url,
|
||
)
|
||
celery_app.conf.update(
|
||
task_serializer="json",
|
||
accept_content=["json"],
|
||
result_serializer="json",
|
||
timezone="Asia/Shanghai",
|
||
task_routes={
|
||
"app.worker.process_file_task": {"queue": "parse"},
|
||
"app.worker.fetch_regulation_source": {"queue": "monitor"},
|
||
"app.worker.send_notifications": {"queue": "push"},
|
||
},
|
||
beat_schedule={
|
||
"daily-regulation-monitor": {
|
||
"task": "app.worker.run_all_monitors",
|
||
"schedule": crontab(hour=2, minute=0),
|
||
},
|
||
},
|
||
)
|
||
|
||
# ── 文件处理任务（解析 + 向量化）────────────────
|
||
|
||
@celery_app.task(name="app.worker.process_file_task", bind=True, max_retries=3)
|
||
def process_file_task(self, file_id: str, task_id: str, workspace_id: str):
|
||
"""解析文档并向量化存入 Milvus"""
|
||
import asyncio
|
||
asyncio.run(_process_file(file_id, task_id, workspace_id))
|
||
|
||
|
||
async def _process_file(file_id: str, task_id: str, workspace_id: str):
|
||
from pathlib import Path
|
||
from sqlalchemy import select
|
||
from .core.deps import AsyncSessionLocal, get_milvus_collection
|
||
from .models.db import File, Task
|
||
from .services.parse import parse_document, chunk_text
|
||
from .services.embed import embed_texts
|
||
|
||
async with AsyncSessionLocal() as db:
|
||
# 查找文件记录
|
||
result = await db.execute(select(File).where(File.id == uuid.UUID(file_id)))
|
||
file_record = result.scalar_one_or_none()
|
||
if not file_record:
|
||
logger.error(f"文件 {file_id} 不存在")
|
||
return
|
||
|
||
task_result = await db.execute(select(Task).where(Task.id == uuid.UUID(task_id)))
|
||
task = task_result.scalar_one_or_none()
|
||
|
||
try:
|
||
# 更新状态
|
||
file_record.status = "parsing"
|
||
if task:
|
||
task.status = "running"
|
||
task.progress = 10
|
||
await db.commit()
|
||
|
||
# Step 1：解析文档
|
||
file_content = Path(file_record.storage_path).read_bytes()
|
||
parse_result = await parse_document(file_content, file_record.original_name)
|
||
markdown = parse_result.get("markdown", "")
|
||
|
||
if not markdown.strip():
|
||
raise ValueError("文档解析结果为空")
|
||
|
||
file_record.status = "parsed"
|
||
if task:
|
||
task.progress = 40
|
||
await db.commit()
|
||
|
||
# Step 2：分块
|
||
chunks = chunk_text(markdown, chunk_size=512, overlap=64)
|
||
logger.info(f"文件 {file_id} 分割为 {len(chunks)} 块")
|
||
|
||
# Step 3：向量化（分批处理）
|
||
batch_size = 16
|
||
col = get_milvus_collection("regulation_chunks")
|
||
|
||
for i in range(0, len(chunks), batch_size):
|
||
batch = chunks[i:i + batch_size]
|
||
texts = [c["content"] for c in batch]
|
||
embed_result = await embed_texts(texts, batch_size=batch_size)
|
||
dense_vecs = embed_result["dense"]
|
||
|
||
entities = [
|
||
[f"{file_id}_{c['idx']}" for c in batch],
|
||
[file_id] * len(batch),
|
||
[workspace_id] * len(batch),
|
||
[c["idx"] for c in batch],
|
||
[c["content"] for c in batch],
|
||
dense_vecs,
|
||
[{"filename": file_record.original_name, "page": c.get("page", 0)} for c in batch],
|
||
]
|
||
col.insert(entities)
|
||
|
||
if task:
|
||
task.progress = 40 + int(60 * (i + batch_size) / len(chunks))
|
||
await db.commit()
|
||
|
||
col.flush()
|
||
|
||
# 完成
|
||
file_record.status = "vectorized"
|
||
if task:
|
||
task.status = "completed"
|
||
task.progress = 100
|
||
task.completed_at = datetime.now(timezone.utc)
|
||
await db.commit()
|
||
logger.info(f"文件 {file_id} 处理完成")
|
||
|
||
except Exception as e:
|
||
logger.error(f"文件 {file_id} 处理失败：{e}")
|
||
file_record.status = "failed"
|
||
file_record.error_msg = str(e)
|
||
if task:
|
||
task.status = "failed"
|
||
task.error_msg = str(e)
|
||
await db.commit()
|
||
raise
|
||
|
||
|
||
# ── 法规监控任务 ────────────────────────────────
|
||
|
||
@celery_app.task(name="app.worker.run_all_monitors")
|
||
def run_all_monitors():
|
||
"""定时触发所有活跃监控源"""
|
||
import asyncio
|
||
asyncio.run(_run_all_monitors())
|
||
|
||
|
||
async def _run_all_monitors():
|
||
from sqlalchemy import select
|
||
from .core.deps import AsyncSessionLocal
|
||
from .models.db import RegulationSource
|
||
|
||
async with AsyncSessionLocal() as db:
|
||
result = await db.execute(
|
||
select(RegulationSource).where(RegulationSource.is_active == True)
|
||
)
|
||
sources = result.scalars().all()
|
||
for source in sources:
|
||
fetch_regulation_source.delay(str(source.id))
|
||
logger.info(f"触发监控源抓取：{source.name}")
|
||
|
||
|
||
@celery_app.task(name="app.worker.fetch_regulation_source", bind=True, max_retries=2)
|
||
def fetch_regulation_source(self, source_id: str):
|
||
import asyncio
|
||
asyncio.run(_fetch_source(source_id))
|
||
|
||
|
||
async def _fetch_source(source_id: str):
|
||
import hashlib
|
||
from sqlalchemy import select
|
||
from .core.deps import AsyncSessionLocal
|
||
from .models.db import RegulationSource, RegulationUpdate
|
||
from .services.monitor import check_source_for_updates
|
||
|
||
async with AsyncSessionLocal() as db:
|
||
result = await db.execute(
|
||
select(RegulationSource).where(RegulationSource.id == uuid.UUID(source_id))
|
||
)
|
||
source = result.scalar_one_or_none()
|
||
if not source:
|
||
return
|
||
|
||
source_dict = {
|
||
"id": str(source.id),
|
||
"name": source.name,
|
||
"url": source.url,
|
||
"last_hash": source.last_hash,
|
||
}
|
||
update_data = await check_source_for_updates(source_dict)
|
||
|
||
if update_data:
|
||
logger.info(f"检测到变更：{source.name}")
|
||
source.last_hash = update_data["new_hash"]
|
||
source.last_fetched_at = datetime.now(timezone.utc)
|
||
|
||
update = RegulationUpdate(
|
||
source_id=uuid.UUID(source_id),
|
||
change_type="updated",
|
||
raw_content=update_data["raw_content"][:50000],
|
||
importance="normal",
|
||
)
|
||
db.add(update)
|
||
await db.commit()
|
||
else:
|
||
source.last_fetched_at = datetime.now(timezone.utc)
|
||
await db.commit()
|
||
|
||
|
||
@celery_app.task(name="app.worker.send_notifications")
|
||
def send_notifications():
|
||
logger.info("推送通知任务执行（待实现）")
|
||
|
||
|
||
# 导出供 FastAPI 使用
|
||
worker = celery_app
|