406 lines
14 KiB
Python
406 lines
14 KiB
Python
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import shutil
|
||
import sqlite3
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
from asrpipe import graph, lexicon, providers, refine, review, revise, vault
|
||
from asrpipe.config import CFG
|
||
from asrpipe.locking import acquire_recording_lock
|
||
from . import db
|
||
from .config import (
|
||
DATA, DRAFTS, FAILED, INBOX, LOGS, MAX_ATTEMPTS,
|
||
POLL_SECONDS, PROCESSED, PROCESSING, PUBLIC_URL, REVISION_QUEUE,
|
||
STABLE_SECONDS, STATE, SUPPORTED, TRANSCRIPTS, ensure_dirs,
|
||
)
|
||
from .notify import send
|
||
|
||
_RECORDING_LOCK = None
|
||
|
||
|
||
def log(message: str) -> None:
|
||
ensure_dirs()
|
||
with (LOGS / "worker.log").open("a", encoding="utf-8") as stream:
|
||
stream.write(time.strftime("[%Y-%m-%d %H:%M:%S] ") + message + "\n")
|
||
|
||
|
||
def sha256(path: Path) -> str:
|
||
h = hashlib.sha256()
|
||
with path.open("rb") as stream:
|
||
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
||
h.update(chunk)
|
||
return h.hexdigest()
|
||
|
||
|
||
def transcript_from_response(payload: dict) -> str:
|
||
"""Backward-compatible helper for qwen3cam JSON responses."""
|
||
return providers.transcript_from_qwen3cam_response(payload)
|
||
|
||
|
||
def transcribe(
|
||
audio: Path,
|
||
job_id: str,
|
||
*,
|
||
provider: str = "",
|
||
model: str = "",
|
||
title: str = "",
|
||
user_context: str = "",
|
||
) -> tuple[Path, providers.TranscribeResult]:
|
||
res = providers.transcribe_audio(
|
||
audio,
|
||
provider=provider or CFG.asr_provider,
|
||
model=model,
|
||
title=title,
|
||
user_context=user_context,
|
||
)
|
||
response_path = TRANSCRIPTS / f"{job_id}.asr.json"
|
||
try:
|
||
response_path.write_text(
|
||
json.dumps(res.raw_response, ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
except Exception:
|
||
pass
|
||
transcript = TRANSCRIPTS / f"{job_id}.txt"
|
||
transcript.write_text(res.transcript_text, encoding="utf-8")
|
||
return transcript, res
|
||
|
||
|
||
def intake_header(payload: dict) -> str:
|
||
blob = json.dumps(payload, ensure_ascii=False).replace("--", "- -")
|
||
return f"<!-- intake: {blob} -->"
|
||
|
||
|
||
def replace_intake(raw: str, header: str) -> str:
|
||
if raw.startswith("<!-- intake:"):
|
||
_first, _sep, rest = raw.partition("\n")
|
||
return header + ("\n" + rest if _sep else "")
|
||
return header + "\n" + raw
|
||
|
||
|
||
def beat() -> None:
|
||
ensure_dirs()
|
||
(STATE / "worker.heartbeat").write_text(str(time.time()))
|
||
|
||
|
||
def _beat_loop(stop: threading.Event) -> None:
|
||
while not stop.wait(15):
|
||
try:
|
||
beat()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def process_audio(source: Path, job: dict) -> None:
|
||
job_id = job["id"]
|
||
attempts = int(job.get("attempts") or 0) + 1
|
||
prov = providers.normalize_provider_name(job.get("asr_provider") or CFG.asr_provider)
|
||
model = str(job.get("asr_model") or "")
|
||
digest = sha256(source)
|
||
prov_label = providers.PROVIDERS.get(prov, {}).get("label", prov)
|
||
db.update_job(
|
||
job_id, "transcribing", f"正在转写 ({prov_label})",
|
||
sha256=digest, attempts=attempts, error_code="",
|
||
asr_provider=prov, asr_model=model,
|
||
)
|
||
transcript, res = transcribe(
|
||
source, job_id,
|
||
provider=prov, model=model,
|
||
title=str(job.get("title") or ""),
|
||
user_context=str(job.get("context") or ""),
|
||
)
|
||
intake = {
|
||
"job_id": job_id,
|
||
"title": job.get("title", ""),
|
||
"context": job.get("context", ""),
|
||
"source": job.get("source", ""),
|
||
"asr_provider": res.provider,
|
||
"asr_model": res.model,
|
||
}
|
||
raw = transcript.read_text(encoding="utf-8")
|
||
raw = replace_intake(raw, intake_header(intake))
|
||
transcript.write_text(raw, encoding="utf-8")
|
||
db.update_job(
|
||
job_id, "refining", f"转写完成 ({res.provider}:{res.model})，正在生成整理稿",
|
||
transcript=str(transcript), asr_provider=res.provider, asr_model=res.model,
|
||
)
|
||
result = refine.refine(
|
||
raw, date=time.strftime("%Y-%m-%d"),
|
||
transcript_path=str(transcript), queue_review=True,
|
||
)
|
||
draft = DRAFTS / f"{job_id}.md"
|
||
vault.write_draft_file(draft, result.draft)
|
||
questions = review.enqueue_draft(str(transcript), draft)
|
||
db.update_job(
|
||
job_id, "review", f"整理稿已生成，{questions} 个问题待确认",
|
||
transcript=str(transcript), draft=str(draft),
|
||
)
|
||
target = PROCESSED / source.name
|
||
if target.exists():
|
||
target = PROCESSED / f"{source.stem}-{job_id[:8]}{source.suffix}"
|
||
shutil.move(str(source), target)
|
||
send("ASR 草稿待确认", job.get("title") or job["original_name"],
|
||
f"{PUBLIC_URL}/jobs/{job_id}/review/")
|
||
|
||
|
||
def rerefine_job(job: dict) -> None:
|
||
"""Re-run LLM refinement on an existing raw transcript without re-running ASR."""
|
||
job_id = job["id"]
|
||
transcript = Path(str(job.get("transcript") or ""))
|
||
if not transcript.is_file():
|
||
raise FileNotFoundError(f"transcript file missing for job {job_id}: {transcript}")
|
||
raw = transcript.read_text(encoding="utf-8", errors="ignore")
|
||
intake = {
|
||
"job_id": job_id,
|
||
"title": job.get("title", ""),
|
||
"context": job.get("context", ""),
|
||
"source": job.get("source", ""),
|
||
"asr_provider": job.get("asr_provider") or CFG.asr_provider,
|
||
"asr_model": job.get("asr_model") or "",
|
||
}
|
||
raw = replace_intake(raw, intake_header(intake))
|
||
transcript.write_text(raw, encoding="utf-8")
|
||
|
||
db.update_job(job_id, "refining", "正在基于现有转写重新生成精修草稿")
|
||
result = refine.refine(
|
||
raw, date=time.strftime("%Y-%m-%d"),
|
||
transcript_path=str(transcript), queue_review=True,
|
||
)
|
||
draft = DRAFTS / f"{job_id}.md"
|
||
vault.write_draft_file(draft, result.draft)
|
||
review.delete_pending_for_job(
|
||
transcript_path=str(transcript),
|
||
draft_path=str(job.get("draft") or draft),
|
||
job_id=job_id,
|
||
)
|
||
questions = review.enqueue_draft(str(transcript), draft)
|
||
db.update_job(
|
||
job_id, "review", f"整理稿已重新生成，{questions} 个问题待确认",
|
||
transcript=str(transcript), draft=str(draft),
|
||
)
|
||
|
||
|
||
def claim(path: Path) -> tuple[Path, dict]:
|
||
job = db.find_by_stored(path.name)
|
||
if not job:
|
||
job_id = uuid.uuid4().hex[:24]
|
||
job = db.create_job(
|
||
job_id, path.name, path.name, title=path.stem,
|
||
size=path.stat().st_size, source="ssh",
|
||
asr_provider=CFG.asr_provider,
|
||
)
|
||
target = PROCESSING / path.name
|
||
if target.exists():
|
||
target = PROCESSING / f"{path.stem}-{job['id'][:8]}{path.suffix}"
|
||
path.replace(target)
|
||
if target.name != job.get("stored_name"):
|
||
db.update_job(job["id"], job["status"], job.get("detail") or "已认领", stored_name=target.name)
|
||
job = db.get_job(job["id"]) or job
|
||
return target, job
|
||
|
||
|
||
def scan_inbox() -> None:
|
||
now = time.time()
|
||
for path in sorted(INBOX.iterdir()):
|
||
if not path.is_file() or path.suffix.lower() not in SUPPORTED:
|
||
continue
|
||
if now - path.stat().st_mtime < STABLE_SECONDS:
|
||
continue
|
||
source, job = claim(path)
|
||
try:
|
||
process_audio(source, job)
|
||
except Exception as exc:
|
||
attempts = int((db.get_job(job["id"]) or {}).get("attempts") or 0)
|
||
is_empty_audio = isinstance(exc, providers.EmptyAudioError)
|
||
retry = (attempts < MAX_ATTEMPTS) and not is_empty_audio
|
||
db.update_job(
|
||
job["id"], "queued" if retry else "failed",
|
||
f"处理失败，等待重试 ({exc})" if retry else f"处理失败：{exc}",
|
||
error_code=type(exc).__name__,
|
||
)
|
||
log(f"job={job['id']} failed type={type(exc).__name__} retry={retry} err={exc}")
|
||
try:
|
||
if source.exists():
|
||
folder = INBOX if retry else FAILED
|
||
target = folder / source.name
|
||
if target.exists():
|
||
target = folder / f"{source.stem}-{job['id'][:8]}{source.suffix}"
|
||
shutil.move(str(source), target)
|
||
if retry:
|
||
old = time.time() - max(STABLE_SECONDS, 60)
|
||
os.utime(target, (old, old))
|
||
except Exception as move_exc:
|
||
log(f"job={job['id']} audio left in processing after failure: {move_exc}")
|
||
|
||
|
||
def _defer_revision(marker: Path, payload: dict, exc: BaseException) -> None:
|
||
attempts = int(payload.get("attempts") or 0) + 1
|
||
payload["attempts"] = attempts
|
||
payload["last_error"] = str(exc)[:300]
|
||
job_id = str(payload.get("job_id") or "")
|
||
if attempts >= MAX_ATTEMPTS:
|
||
marker.unlink(missing_ok=True)
|
||
if job_id:
|
||
db.update_job(job_id, "failed", f"修订失败：{exc}", error_code=type(exc).__name__)
|
||
log(f"revision marker={marker.name} gave up after {attempts}: {exc}")
|
||
return
|
||
tmp = marker.with_name(f".{marker.name}.{os.getpid()}.tmp")
|
||
tmp.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
|
||
tmp.replace(marker)
|
||
if job_id:
|
||
db.update_job(
|
||
job_id, "review",
|
||
f"修订失败，将重试 ({attempts}/{MAX_ATTEMPTS})：{exc}",
|
||
error_code=type(exc).__name__,
|
||
)
|
||
log(f"revision marker={marker.name} retry {attempts}/{MAX_ATTEMPTS}: {exc}")
|
||
|
||
|
||
def process_revisions() -> None:
|
||
for marker in sorted(REVISION_QUEUE.glob("*.json")):
|
||
if marker.name.startswith("."):
|
||
continue
|
||
try:
|
||
payload = json.loads(marker.read_text(encoding="utf-8"))
|
||
except Exception as exc:
|
||
log(f"revision marker={marker.name} unreadable: {exc}")
|
||
_defer_revision(marker, {}, exc)
|
||
continue
|
||
try:
|
||
action = str(payload.get("action") or "revise")
|
||
job_id = str(payload.get("job_id") or "")
|
||
if action == "rerefine":
|
||
job = db.get_job(job_id)
|
||
if job:
|
||
rerefine_job(job)
|
||
marker.unlink(missing_ok=True)
|
||
continue
|
||
draft = Path(payload["draft"])
|
||
db.update_job(job_id, "publishing", "正在应用人工确认")
|
||
revise.revise(draft)
|
||
current = db.get_job(job_id) or {}
|
||
if current.get("status") == "complete":
|
||
send("ASR 正式记录已生成",
|
||
current.get("title") or current.get("original_name", ""),
|
||
f"{PUBLIC_URL}/jobs/{job_id}/")
|
||
elif current.get("status") != "review":
|
||
db.update_job(
|
||
job_id, "review", "定稿预览已生成，等待发布确认",
|
||
draft=str(draft),
|
||
)
|
||
marker.unlink(missing_ok=True)
|
||
except Exception as exc:
|
||
_defer_revision(marker, payload, exc)
|
||
|
||
|
||
def _park_audio(path: Path) -> Path:
|
||
dest = INBOX / path.name
|
||
if dest.exists():
|
||
dest = INBOX / f"{path.stem}-recovered{path.suffix}"
|
||
shutil.move(str(path), dest)
|
||
old = time.time() - max(STABLE_SECONDS, 60)
|
||
os.utime(dest, (old, old))
|
||
return dest
|
||
|
||
|
||
def recover() -> None:
|
||
"""Return audio left in processing/ to the inbox and requeue its job."""
|
||
if not PROCESSING.exists():
|
||
return
|
||
files = [p for p in PROCESSING.iterdir() if p.is_file()]
|
||
jobs = db.pending_recovery()
|
||
claimed: set[str] = set()
|
||
for path in files:
|
||
match = None
|
||
for job in jobs:
|
||
if job["id"] in claimed:
|
||
continue
|
||
stored = str(job.get("stored_name") or "")
|
||
if path.name == stored or (job["id"][:8] and job["id"][:8] in path.name):
|
||
match = job
|
||
break
|
||
dest = _park_audio(path)
|
||
if match:
|
||
claimed.add(match["id"])
|
||
db.update_job(
|
||
match["id"], "queued", "worker 重启后重新排队",
|
||
stored_name=dest.name, error_code="",
|
||
)
|
||
log(f"recovered job={match['id']} audio={dest.name}")
|
||
else:
|
||
log(f"recovered orphan audio={dest.name}")
|
||
|
||
|
||
def sync_vault_knowledge() -> None:
|
||
try:
|
||
vault.pull()
|
||
lexicon.sync_from_vault()
|
||
graph.sync()
|
||
except Exception as exc:
|
||
log(f"vault knowledge sync skipped: {exc}")
|
||
|
||
|
||
def backup_databases() -> None:
|
||
backup_dir = DATA / "backups"
|
||
backup_dir.mkdir(parents=True, exist_ok=True)
|
||
for name in ("intake.db", "lexicon.db"):
|
||
source = STATE / name
|
||
if not source.exists():
|
||
continue
|
||
target = backup_dir / f"{name}.latest"
|
||
temp = backup_dir / f".{name}.{os.getpid()}.tmp"
|
||
with sqlite3.connect(source) as src, sqlite3.connect(temp) as dst:
|
||
src.backup(dst)
|
||
temp.replace(target)
|
||
|
||
|
||
def run_loop() -> None:
|
||
os.environ.setdefault("ASR_ROLE", "worker")
|
||
os.environ.setdefault("ASR_REVISION_QUEUE_DIR", str(REVISION_QUEUE))
|
||
ensure_dirs()
|
||
global _RECORDING_LOCK
|
||
lock_path = Path(os.getenv("ASR_CONSUMER_LOCK", str(STATE / "consumer.lock")))
|
||
_RECORDING_LOCK = acquire_recording_lock(lock_path)
|
||
recover()
|
||
sync_vault_knowledge()
|
||
beat()
|
||
stop_beat = threading.Event()
|
||
threading.Thread(target=_beat_loop, args=(stop_beat,), name="asr-heartbeat", daemon=True).start()
|
||
last_backup = 0.0
|
||
last_vault_sync = time.time()
|
||
while True:
|
||
try:
|
||
scan_inbox()
|
||
except Exception as exc:
|
||
log(f"scan_inbox failed: {exc}")
|
||
try:
|
||
process_revisions()
|
||
except Exception as exc:
|
||
log(f"process_revisions failed: {exc}")
|
||
now = time.time()
|
||
if now - last_vault_sync >= 900:
|
||
sync_vault_knowledge()
|
||
last_vault_sync = now
|
||
if now - last_backup >= 3600:
|
||
try:
|
||
backup_databases()
|
||
except Exception as exc:
|
||
log(f"backup failed: {exc}")
|
||
last_backup = now
|
||
beat()
|
||
time.sleep(POLL_SECONDS)
|
||
|
||
|
||
def main() -> None:
|
||
run_loop()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|