“ydy0615”
257 lines
8.3 KiB
Python
257 lines
8.3 KiB
Python
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)
|