Migrate backend jobs to Redis Streams

2026-06-06 15:44:00 +08:00
parent 2c7a02f587
commit 81f711ef0b
69 changed files with 2709 additions and 6291 deletions
@@ -0,0 +1,256 @@
+import asyncio
+import os
+import re
+from contextlib import suppress
+from typing import Any, Callable, Awaitable
+
+import markitdown
+
+from llm import call_ollama, call_vlm_ocr, stream_ollama_events
+from prompt import (
+    build_completion_prompts,
+    build_pro_completion_prompts,
+    prepare_prompt_context,
+)
+
+try:  # pragma: no cover - optional heavy dependency path
+    from tts_asr import generate_asr_response, generate_tts_response
+except Exception:  # pragma: no cover
+    generate_tts_response = None
+    generate_asr_response = None
+
+
+IMAGE_MARKDOWN_RE = re.compile(r"!\[[^\]]*]\([^)]+\)")
+IMAGE_HTML_RE = re.compile(r"<img\b[^>]*>", re.IGNORECASE)
+ALLOWED_CONVERT_EXTENSIONS = {".txt", ".docx", ".pptx", ".pdf"}
+_markitdown_instance = None
+
+
+def _get_markitdown():
+    global _markitdown_instance
+    if _markitdown_instance is None:
+        _markitdown_instance = markitdown.MarkItDown()
+    return _markitdown_instance
+
+
+def _safe_unlink(path: str | None) -> None:
+    if not path:
+        return
+    with suppress(FileNotFoundError):
+        os.unlink(path)
+
+
+def _sanitize_converted_markdown(text: str) -> str:
+    value = (text or "").replace("\r\n", "\n").replace("\r", "\n")
+    value = IMAGE_MARKDOWN_RE.sub("", value)
+    value = IMAGE_HTML_RE.sub("", value)
+    value = re.sub(r"\n{3,}", "\n\n", value)
+    return value.strip()
+
+
+def sanitize_inline_completion_content(text: str, prefill: str = "") -> str:
+    value = (text or "").strip()
+    if not value:
+        return ""
+
+    fim_middle = value.rfind("<|fim_middle|>")
+    if fim_middle >= 0:
+        value = value[fim_middle + len("<|fim_middle|>") :]
+
+    end_index = value.find("<|end|>")
+    if end_index >= 0:
+        value = value[:end_index]
+
+    quoted = re.findall(r'"([^"]+)"', value)
+    if quoted:
+        value = quoted[-1]
+
+    marker_index = max(value.rfind("|fim_middle|>"), value.rfind("<|start|>assistant"))
+    if marker_index >= 0:
+        tail = value.split(">")[-1]
+        if tail:
+            value = tail
+
+    value = value.strip()
+    if prefill and value.startswith(prefill):
+        value = value[len(prefill) :]
+
+    return value.strip()
+
+
+async def completion_handler(
+    payload: dict[str, Any],
+    emit: Callable[[str, dict[str, Any]], Awaitable[None]],
+    is_cancelled: Callable[[], bool],
+) -> dict[str, Any]:
+    req = payload["request"]
+    system_prompt, user_prompt, prefill = build_completion_prompts(
+        req["prefix"],
+        req["suffix"],
+        req.get("languageId", "markdown"),
+        location=payload.get("location", ""),
+        thinking_level=req.get("model_thinking", "low"),
+        preferences=req.get("user_preferences"),
+    )
+
+    result = await call_ollama(
+        user_prompt,
+        system_prompt=system_prompt,
+        tag=f'{payload["request_id"][:8]}-completion',
+        temperature=float(req.get("temperature", 0.7)),
+        thinking=req.get("model_thinking") if req.get("model_thinking") != "none" else None,
+        model=req.get("model"),
+        prefill=prefill or None,
+    )
+    content = sanitize_inline_completion_content(result.get("content") or "", prefill=prefill or "")
+    if is_cancelled():
+        raise asyncio.CancelledError()
+    await emit("result", {"content": content})
+    return {"content": content, "request_id": payload["request_id"]}
+
+
+async def pro_completion_handler(
+    payload: dict[str, Any],
+    emit: Callable[[str, dict[str, Any]], Awaitable[None]],
+    is_cancelled: Callable[[], bool],
+) -> dict[str, Any]:
+    req = payload["request"]
+    system_prompt, user_prompt = build_pro_completion_prompts(
+        prefix=req["prefix"],
+        suffix=req["suffix"],
+        instruction=req.get("instruction", ""),
+        language_id=req.get("languageId", "markdown"),
+        location=payload.get("location", ""),
+        pro_thinking_level=req.get("pro_thinking", "medium"),
+        preferences=req.get("user_preferences"),
+    )
+    chunks: list[str] = []
+    async for event_type, delta in stream_ollama_events(
+        user_prompt,
+        system_prompt=system_prompt,
+        tag=f'{payload["request_id"][:8]}-pro',
+        temperature=0.7,
+        thinking=req.get("pro_thinking", "medium"),
+        use_pro_model=True,
+        enable_thinking=True,
+    ):
+        if is_cancelled():
+            raise asyncio.CancelledError()
+        if event_type == "thinking":
+            await emit("progress", {"phase": "thinking"})
+            continue
+        if delta:
+            chunks.append(delta)
+            await emit("result", {"delta": delta})
+    content = "".join(chunks)
+    return {"content": content, "request_id": payload["request_id"]}
+
+
+async def compress_handler(
+    payload: dict[str, Any],
+    emit: Callable[[str, dict[str, Any]], Awaitable[None]],
+    is_cancelled: Callable[[], bool],
+) -> dict[str, Any]:
+    content = payload["content"]
+    doc_type = payload.get("docType", "txt")
+    system_prompt = (
+        f"你是一个专业的文档摘要助手。请将以下 {doc_type} 类型文档内容进行精简压缩，"
+        "保留核心信息和关键要点，去除冗余和啰嗦的表述。"
+        "请直接输出压缩后的内容，不要添加任何解释性文字。"
+    )
+    result = await call_ollama(
+        content,
+        system_prompt=system_prompt,
+        tag=f'{payload["request_id"][:8]}-compress',
+    )
+    if is_cancelled():
+        raise asyncio.CancelledError()
+    compressed = result.get("content") or ""
+    await emit("result", {"content": compressed})
+    return {"content": compressed, "request_id": payload["request_id"]}
+
+
+async def ocr_handler(
+    payload: dict[str, Any],
+    emit: Callable[[str, dict[str, Any]], Awaitable[None]],
+    is_cancelled: Callable[[], bool],
+) -> dict[str, Any]:
+    path = payload["input_path"]
+    try:
+        with open(path, "rb") as handle:
+            image_bytes = handle.read()
+        text = await call_vlm_ocr(image_bytes, payload.get("language", "auto"))
+        if is_cancelled():
+            raise asyncio.CancelledError()
+        await emit("result", {"text": text})
+        return {"text": text, "filename": payload.get("filename", "image.jpg")}
+    finally:
+        _safe_unlink(path)
+
+
+async def convert_handler(
+    payload: dict[str, Any],
+    emit: Callable[[str, dict[str, Any]], Awaitable[None]],
+    is_cancelled: Callable[[], bool],
+) -> dict[str, Any]:
+    path = payload["input_path"]
+    filename = payload.get("filename", "document")
+    ext = os.path.splitext(filename)[1].lower()
+    if ext not in ALLOWED_CONVERT_EXTENSIONS:
+        _safe_unlink(path)
+        raise ValueError("仅支持 txt、docx、pptx、pdf 格式")
+    try:
+        if ext == ".txt":
+            with open(path, "rb") as handle:
+                markdown = _sanitize_converted_markdown(handle.read().decode("utf-8", errors="ignore"))
+        else:
+            md = _get_markitdown()
+            result = await asyncio.to_thread(md.convert, path)
+            markdown = _sanitize_converted_markdown(result.text_content)
+        if is_cancelled():
+            raise asyncio.CancelledError()
+        await emit("result", {"markdown": markdown})
+        return {"markdown": markdown, "filename": filename}
+    finally:
+        _safe_unlink(path)
+
+
+async def tts_handler(
+    payload: dict[str, Any],
+    emit: Callable[[str, dict[str, Any]], Awaitable[None]],
+    is_cancelled: Callable[[], bool],
+) -> dict[str, Any]:
+    if generate_tts_response is None:
+        raise RuntimeError("TTS 功能当前不可用")
+    response = await generate_tts_response(
+        text=payload["text"],
+        instruct=payload.get("instruct", ""),
+        speaker=payload.get("speaker", "Vivian"),
+        output_format=payload.get("format", "wav"),
+    )
+    if is_cancelled():
+        raise asyncio.CancelledError()
+    result = response.dict()
+    await emit("result", result)
+    return result
+
+
+async def asr_handler(
+    payload: dict[str, Any],
+    emit: Callable[[str, dict[str, Any]], Awaitable[None]],
+    is_cancelled: Callable[[], bool],
+) -> dict[str, Any]:
+    if generate_asr_response is None:
+        raise RuntimeError("ASR 功能当前不可用")
+    path = payload["input_path"]
+    try:
+        with open(path, "rb") as handle:
+            audio_bytes = handle.read()
+        response = await generate_asr_response(audio_bytes, payload.get("language", "zh-CN"))
+        if is_cancelled():
+            raise asyncio.CancelledError()
+        result = response.dict()
+        await emit("result", result)
+        return result
+    finally:
+        _safe_unlink(path)