139 lines
5.1 KiB
Python
139 lines
5.1 KiB
Python
import asyncio
|
|
import json
|
|
from typing import Any, Dict, List
|
|
|
|
from config import WIKI_CONFIG
|
|
from wiki_engine import db
|
|
|
|
|
|
_WORKER_TASKS: List[asyncio.Task] = []
|
|
_STARTED = False
|
|
|
|
|
|
async def start_wiki_queue_workers() -> None:
|
|
"""Start durable Wiki queue workers for background sync jobs."""
|
|
global _STARTED
|
|
if _STARTED or not WIKI_CONFIG.get("queue_enabled", True):
|
|
return
|
|
_STARTED = True
|
|
worker_count = _bounded_int(WIKI_CONFIG.get("queue_worker_concurrency"), 2, 1, 8)
|
|
for index in range(worker_count):
|
|
_WORKER_TASKS.append(asyncio.create_task(_worker_loop(index + 1)))
|
|
print(f"[wiki_worker] started {worker_count} worker(s)")
|
|
|
|
|
|
async def stop_wiki_queue_workers() -> None:
|
|
for task in _WORKER_TASKS:
|
|
task.cancel()
|
|
if _WORKER_TASKS:
|
|
await asyncio.gather(*_WORKER_TASKS, return_exceptions=True)
|
|
_WORKER_TASKS.clear()
|
|
|
|
|
|
async def _worker_loop(worker_id: int) -> None:
|
|
idle_sleep = 1.0
|
|
while True:
|
|
try:
|
|
ops = await db.claim_wiki_file_ops(
|
|
limit=1,
|
|
stale_seconds=_bounded_int(WIKI_CONFIG.get("queue_claim_stale_seconds"), 1800, 60, 24 * 3600),
|
|
per_kb_concurrency=_bounded_int(WIKI_CONFIG.get("queue_per_kb_concurrency"), 1, 1, 8),
|
|
)
|
|
if not ops:
|
|
finalize_ops = await db.claim_wiki_finalize_ops(limit=1)
|
|
if finalize_ops:
|
|
for finalize_op in finalize_ops:
|
|
await _process_finalize_op(worker_id, finalize_op)
|
|
continue
|
|
await asyncio.sleep(idle_sleep)
|
|
continue
|
|
for op in ops:
|
|
await _process_op(worker_id, op)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
print(f"[wiki_worker] worker {worker_id} loop failed: {exc}")
|
|
await asyncio.sleep(3.0)
|
|
|
|
|
|
async def _process_op(worker_id: int, op: Dict[str, Any]) -> None:
|
|
payload = _decode_payload(op.get("payload"))
|
|
job_id = str(payload.get("job_id") or op.get("job_id") or "")
|
|
kb_id = str(payload.get("kb_id") or op.get("kb_id") or "")
|
|
file_id = str(payload.get("file_id") or op.get("file_id") or "")
|
|
file_info = payload.get("file_info") if isinstance(payload.get("file_info"), dict) else {}
|
|
use_llm = bool(payload.get("use_llm"))
|
|
force = bool(payload.get("force"))
|
|
print(f"[wiki_worker] worker {worker_id} sync_file file_id={file_id} kb_id={kb_id} job_id={job_id}")
|
|
|
|
try:
|
|
from wiki_engine import service
|
|
|
|
result = await service.sync_file(
|
|
file_id=file_id,
|
|
kb_id=kb_id,
|
|
file_info=file_info,
|
|
use_llm=use_llm,
|
|
job_id=job_id,
|
|
force=force,
|
|
)
|
|
await db.complete_wiki_op(int(op["id"]))
|
|
await db.increment_job_progress(
|
|
job_id,
|
|
processed_delta=1,
|
|
slice_delta=int(result.get("slice_count") or 0),
|
|
skipped_delta=1 if result.get("skipped") else 0,
|
|
page_delta=int(result.get("page_count") or 0) if not result.get("skipped") else 0,
|
|
stage="queued_finalize" if kb_id else "file_done",
|
|
)
|
|
if job_id and kb_id:
|
|
job = await db.get_job(job_id)
|
|
if job and int(job.get("processed_files") or 0) >= int(job.get("total_files") or 0):
|
|
await service.enqueue_finalize_job(kb_id=kb_id, job_id=job_id)
|
|
except Exception as exc:
|
|
terminal = await db.fail_wiki_op(
|
|
op,
|
|
error=str(exc),
|
|
max_retries=_bounded_int(WIKI_CONFIG.get("queue_max_retries"), 3, 1, 10),
|
|
)
|
|
if terminal and job_id:
|
|
await db.increment_job_progress(job_id, processed_delta=1, slice_delta=0, error=f"{file_id}: {exc}")
|
|
print(f"[wiki_worker] worker {worker_id} sync_file failed file_id={file_id}: {exc}")
|
|
|
|
|
|
async def _process_finalize_op(worker_id: int, op: Dict[str, Any]) -> None:
|
|
payload = _decode_payload(op.get("payload"))
|
|
job_id = str(payload.get("job_id") or op.get("job_id") or "")
|
|
kb_id = str(payload.get("kb_id") or op.get("kb_id") or "")
|
|
print(f"[wiki_worker] worker {worker_id} finalize_kb kb_id={kb_id} job_id={job_id}")
|
|
try:
|
|
from wiki_engine import service
|
|
|
|
await service.finalize_kb(kb_id=kb_id, job_id=job_id)
|
|
await db.complete_wiki_finalize_op(int(op["id"]))
|
|
except Exception as exc:
|
|
await db.fail_wiki_finalize_op(int(op["id"]), str(exc))
|
|
if job_id:
|
|
await db.update_job(job_id, stage="finalize_failed", error=str(exc))
|
|
print(f"[wiki_worker] worker {worker_id} finalize failed kb_id={kb_id}: {exc}")
|
|
|
|
|
|
def _decode_payload(value: Any) -> Dict[str, Any]:
|
|
if isinstance(value, dict):
|
|
return value
|
|
if isinstance(value, str) and value.strip():
|
|
try:
|
|
parsed = json.loads(value)
|
|
return parsed if isinstance(parsed, dict) else {}
|
|
except Exception:
|
|
return {}
|
|
return {}
|
|
|
|
|
|
def _bounded_int(value: Any, default: int, minimum: int, maximum: int) -> int:
|
|
try:
|
|
number = int(value)
|
|
except Exception:
|
|
number = default
|
|
return max(minimum, min(maximum, number))
|