Migrate backend jobs to Redis Streams
This commit is contained in:
+35
-5
@@ -15,12 +15,42 @@ VLM_MODEL=qwen3-vl:30b
|
||||
# API key for the FastAPI app (change in production)
|
||||
API_KEY=your-secret-key-here
|
||||
|
||||
# PRO completion timeout (seconds)
|
||||
PRO_COMPLETION_TIMEOUT=1200
|
||||
# Job backend
|
||||
JOB_BACKEND=redis
|
||||
REDIS_URL=redis://localhost:6379/0
|
||||
JOB_REDIS_PREFIX=llmtext:jobs
|
||||
JOB_CONSUMER_NAME=
|
||||
JOB_SHARED_TEMP_DIR=/tmp/llm-in-text-jobs
|
||||
JOB_STATE_TTL_SECONDS=600
|
||||
JOB_EVENT_TTL_SECONDS=600
|
||||
JOB_EVENT_STREAM_MAXLEN=512
|
||||
JOB_CANCEL_POLL_SECONDS=0.5
|
||||
JOB_BUSY_NORMAL_THRESHOLD=0.25
|
||||
JOB_BUSY_HIGH_THRESHOLD=0.75
|
||||
JOB_BUSY_FULL_THRESHOLD=1.0
|
||||
|
||||
# Concurrency limits
|
||||
STANDARD_CONCURRENCY_LIMIT=5
|
||||
PRO_CONCURRENCY_LIMIT=20
|
||||
# Per-queue concurrency and capacity
|
||||
JOB_COMPLETION_CONCURRENCY=2
|
||||
JOB_COMPLETION_MAX_QUEUE=16
|
||||
JOB_PRO_COMPLETION_CONCURRENCY=1
|
||||
JOB_PRO_COMPLETION_MAX_QUEUE=8
|
||||
JOB_COMPRESS_CONCURRENCY=1
|
||||
JOB_COMPRESS_MAX_QUEUE=8
|
||||
JOB_OCR_CONCURRENCY=1
|
||||
JOB_OCR_MAX_QUEUE=8
|
||||
JOB_CONVERT_CONCURRENCY=1
|
||||
JOB_CONVERT_MAX_QUEUE=8
|
||||
JOB_TTS_CONCURRENCY=1
|
||||
JOB_TTS_MAX_QUEUE=4
|
||||
JOB_ASR_CONCURRENCY=1
|
||||
JOB_ASR_MAX_QUEUE=4
|
||||
|
||||
# Timeouts (seconds)
|
||||
LLM_COMPLETION_TIMEOUT=600
|
||||
LLM_OCR_TIMEOUT=600
|
||||
|
||||
# Compression limit
|
||||
DOC_COMPRESS_CONTEXT_LIMIT=128000
|
||||
|
||||
# Legacy fallback: if LLM_BASE_URL is not set, OLLAMA_HOST will be auto-converted to /v1/ path
|
||||
#OLLAMA_HOST=http://localhost:11434
|
||||
|
||||
@@ -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)
|
||||
@@ -0,0 +1,704 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncIterator, Awaitable, Callable, Optional
|
||||
|
||||
logger = logging.getLogger("job_system")
|
||||
|
||||
try: # pragma: no cover - optional dependency in tests
|
||||
from redis import asyncio as redis_asyncio # type: ignore
|
||||
except Exception: # pragma: no cover - optional dependency in tests
|
||||
redis_asyncio = None
|
||||
|
||||
|
||||
TERMINAL_STATUSES = {"completed", "failed", "cancelled"}
|
||||
TERMINAL_EVENTS = {"done", "error", "cancelled"}
|
||||
|
||||
JOB_TYPES = (
|
||||
"completion",
|
||||
"pro_completion",
|
||||
"compress",
|
||||
"ocr",
|
||||
"convert",
|
||||
"tts",
|
||||
"asr",
|
||||
)
|
||||
|
||||
DEFAULT_CONCURRENCY = {
|
||||
"completion": 2,
|
||||
"pro_completion": 1,
|
||||
"compress": 1,
|
||||
"ocr": 1,
|
||||
"convert": 1,
|
||||
"tts": 1,
|
||||
"asr": 1,
|
||||
}
|
||||
|
||||
DEFAULT_QUEUE_SIZE = {
|
||||
"completion": 16,
|
||||
"pro_completion": 8,
|
||||
"compress": 8,
|
||||
"ocr": 8,
|
||||
"convert": 8,
|
||||
"tts": 4,
|
||||
"asr": 4,
|
||||
}
|
||||
|
||||
|
||||
class JobSystemError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class QueueFullError(JobSystemError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueueConfig:
|
||||
job_type: str
|
||||
concurrency: int
|
||||
max_queue: int
|
||||
|
||||
|
||||
Handler = Callable[[dict[str, Any], Callable[[str, dict[str, Any]], Awaitable[None]], Callable[[], bool]], Awaitable[dict[str, Any]]]
|
||||
|
||||
|
||||
def _bool_env(name: str, default: bool) -> bool:
|
||||
value = os.getenv(name)
|
||||
if value is None:
|
||||
return default
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _int_env(name: str, default: int) -> int:
|
||||
try:
|
||||
return max(1, int(os.getenv(name, str(default))))
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _float_env(name: str, default: float) -> float:
|
||||
try:
|
||||
return float(os.getenv(name, str(default)))
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _now_ms() -> int:
|
||||
return int(time.time() * 1000)
|
||||
|
||||
|
||||
def _busy_level(ratio: float) -> str:
|
||||
if ratio >= _float_env("JOB_BUSY_FULL_THRESHOLD", 1.0):
|
||||
return "full"
|
||||
if ratio >= _float_env("JOB_BUSY_HIGH_THRESHOLD", 0.75):
|
||||
return "busy"
|
||||
if ratio >= _float_env("JOB_BUSY_NORMAL_THRESHOLD", 0.25):
|
||||
return "normal"
|
||||
return "idle"
|
||||
|
||||
|
||||
def _json_dumps(value: Any) -> str:
|
||||
return json.dumps(value, ensure_ascii=False)
|
||||
|
||||
|
||||
def _json_loads(value: str | bytes | None, default: Any = None) -> Any:
|
||||
if value is None:
|
||||
return default
|
||||
if isinstance(value, bytes):
|
||||
value = value.decode("utf-8")
|
||||
if not value:
|
||||
return default
|
||||
return json.loads(value)
|
||||
|
||||
|
||||
def _queue_config(job_type: str) -> QueueConfig:
|
||||
upper = job_type.upper()
|
||||
concurrency = _int_env(f"JOB_{upper}_CONCURRENCY", DEFAULT_CONCURRENCY[job_type])
|
||||
max_queue = _int_env(f"JOB_{upper}_MAX_QUEUE", DEFAULT_QUEUE_SIZE[job_type])
|
||||
return QueueConfig(job_type=job_type, concurrency=concurrency, max_queue=max_queue)
|
||||
|
||||
|
||||
def get_job_backend_name() -> str:
|
||||
value = (os.getenv("JOB_BACKEND") or "").strip().lower()
|
||||
if value:
|
||||
return value
|
||||
if redis_asyncio is not None:
|
||||
return "redis"
|
||||
return "memory"
|
||||
|
||||
|
||||
def _shared_temp_dir() -> Path:
|
||||
path = Path(os.getenv("JOB_SHARED_TEMP_DIR", tempfile.gettempdir()) or tempfile.gettempdir())
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def persist_temp_input(raw_bytes: bytes, suffix: str) -> str:
|
||||
directory = _shared_temp_dir()
|
||||
fd, path = tempfile.mkstemp(prefix="job-input-", suffix=suffix, dir=directory)
|
||||
os.close(fd)
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(raw_bytes)
|
||||
return path
|
||||
|
||||
|
||||
async def _maybe_await(value: Any) -> Any:
|
||||
if inspect.isawaitable(value):
|
||||
return await value
|
||||
return value
|
||||
|
||||
|
||||
class BaseJobManager:
|
||||
def __init__(self) -> None:
|
||||
self.handlers: dict[str, Handler] = {}
|
||||
|
||||
def register_handler(self, job_type: str, handler: Handler) -> None:
|
||||
self.handlers[job_type] = handler
|
||||
|
||||
async def submit(self, job_type: str, payload: dict[str, Any], request_id: str | None = None) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
async def cancel(self, job_id: str, reason: str = "abort") -> dict[str, Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
async def get_status(self, job_id: str) -> dict[str, Any] | None:
|
||||
raise NotImplementedError
|
||||
|
||||
async def stream_events(self, job_id: str) -> AsyncIterator[dict[str, Any]]:
|
||||
raise NotImplementedError
|
||||
|
||||
async def close(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class InMemoryJobManager(BaseJobManager):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.jobs: dict[str, dict[str, Any]] = {}
|
||||
self.event_history: dict[str, list[dict[str, Any]]] = {}
|
||||
self.subscribers: dict[str, list[asyncio.Queue]] = {}
|
||||
self.queues = {job_type: asyncio.Queue() for job_type in JOB_TYPES}
|
||||
self.semaphores = {job_type: asyncio.Semaphore(_queue_config(job_type).concurrency) for job_type in JOB_TYPES}
|
||||
self.queue_counts = {job_type: 0 for job_type in JOB_TYPES}
|
||||
self.running_counts = {job_type: 0 for job_type in JOB_TYPES}
|
||||
self.running_tasks: dict[str, asyncio.Task] = {}
|
||||
self.worker_tasks: list[asyncio.Task] = []
|
||||
self.started = False
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
async def _ensure_started(self) -> None:
|
||||
if self.started:
|
||||
return
|
||||
self.started = True
|
||||
for job_type in JOB_TYPES:
|
||||
self.worker_tasks.append(asyncio.create_task(self._worker_loop(job_type)))
|
||||
|
||||
def _metrics(self, job_type: str) -> dict[str, Any]:
|
||||
config = _queue_config(job_type)
|
||||
queued = self.queue_counts[job_type]
|
||||
running = self.running_counts[job_type]
|
||||
capacity = max(config.max_queue + config.concurrency, 1)
|
||||
ratio = min((queued + running) / capacity, 1.0)
|
||||
return {
|
||||
"queue_position": queued if queued > 0 else 0,
|
||||
"queued_count": queued,
|
||||
"running_count": running,
|
||||
"concurrency_limit": config.concurrency,
|
||||
"max_queue": config.max_queue,
|
||||
"busy_ratio": round(ratio, 4),
|
||||
"busy_level": _busy_level(ratio),
|
||||
}
|
||||
|
||||
async def _publish(self, job_id: str, event: str, data: dict[str, Any]) -> None:
|
||||
event_payload = {"event": event, **data}
|
||||
self.event_history.setdefault(job_id, []).append(event_payload)
|
||||
for queue in self.subscribers.get(job_id, []):
|
||||
await queue.put(event_payload)
|
||||
|
||||
async def submit(self, job_type: str, payload: dict[str, Any], request_id: str | None = None) -> str:
|
||||
await self._ensure_started()
|
||||
if job_type not in self.handlers:
|
||||
raise JobSystemError(f"missing handler for job type: {job_type}")
|
||||
|
||||
async with self.lock:
|
||||
config = _queue_config(job_type)
|
||||
if self.queue_counts[job_type] >= config.max_queue:
|
||||
raise QueueFullError(f"{job_type} queue is full")
|
||||
job_id = request_id or str(uuid.uuid4())
|
||||
self.jobs[job_id] = {
|
||||
"job_id": job_id,
|
||||
"request_id": job_id,
|
||||
"job_type": job_type,
|
||||
"status": "queued",
|
||||
"payload": payload,
|
||||
"result": None,
|
||||
"error": "",
|
||||
"cancel_requested": False,
|
||||
"created_at": _now_ms(),
|
||||
"updated_at": _now_ms(),
|
||||
}
|
||||
self.event_history[job_id] = []
|
||||
self.queue_counts[job_type] += 1
|
||||
metrics = self._metrics(job_type)
|
||||
|
||||
await self._publish(job_id, "queued", {"job_id": job_id, "type": job_type, "status": "queued", **metrics})
|
||||
await self.queues[job_type].put(job_id)
|
||||
return job_id
|
||||
|
||||
async def cancel(self, job_id: str, reason: str = "abort") -> dict[str, Any]:
|
||||
async with self.lock:
|
||||
job = self.jobs.get(job_id)
|
||||
if not job:
|
||||
return {"cancelled": False, "status": "not_found"}
|
||||
if job["status"] in TERMINAL_STATUSES:
|
||||
return {"cancelled": False, "status": job["status"]}
|
||||
job["cancel_requested"] = True
|
||||
job["updated_at"] = _now_ms()
|
||||
task = self.running_tasks.get(job_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
if job["status"] == "queued":
|
||||
job["status"] = "cancelled"
|
||||
self.queue_counts[job["job_type"]] = max(0, self.queue_counts[job["job_type"]] - 1)
|
||||
metrics = self._metrics(job["job_type"])
|
||||
else:
|
||||
job["status"] = "cancelled"
|
||||
metrics = self._metrics(job["job_type"])
|
||||
|
||||
await self._publish(job_id, "cancelled", {"job_id": job_id, "type": job["job_type"], "status": "cancelled", "reason": reason, **metrics})
|
||||
return {"cancelled": True, "status": "ok"}
|
||||
|
||||
async def get_status(self, job_id: str) -> dict[str, Any] | None:
|
||||
job = self.jobs.get(job_id)
|
||||
if not job:
|
||||
return None
|
||||
metrics = self._metrics(job["job_type"])
|
||||
return {
|
||||
"job_id": job_id,
|
||||
"request_id": job["request_id"],
|
||||
"type": job["job_type"],
|
||||
"status": job["status"],
|
||||
"result": job["result"],
|
||||
"error": job["error"],
|
||||
**metrics,
|
||||
}
|
||||
|
||||
async def stream_events(self, job_id: str) -> AsyncIterator[dict[str, Any]]:
|
||||
queue: asyncio.Queue = asyncio.Queue()
|
||||
history = list(self.event_history.get(job_id, []))
|
||||
for item in history:
|
||||
yield item
|
||||
self.subscribers.setdefault(job_id, []).append(queue)
|
||||
try:
|
||||
while True:
|
||||
event = await queue.get()
|
||||
yield event
|
||||
if event["event"] in TERMINAL_EVENTS:
|
||||
break
|
||||
finally:
|
||||
with suppress(ValueError):
|
||||
self.subscribers.get(job_id, []).remove(queue)
|
||||
|
||||
async def _worker_loop(self, job_type: str) -> None:
|
||||
queue = self.queues[job_type]
|
||||
sem = self.semaphores[job_type]
|
||||
while True:
|
||||
job_id = await queue.get()
|
||||
async with self.lock:
|
||||
job = self.jobs.get(job_id)
|
||||
if not job or job["status"] == "cancelled":
|
||||
continue
|
||||
await sem.acquire()
|
||||
task = asyncio.create_task(self._run_job(job_id))
|
||||
self.running_tasks[job_id] = task
|
||||
|
||||
async def _run_job(self, job_id: str) -> None:
|
||||
job = self.jobs[job_id]
|
||||
job_type = job["job_type"]
|
||||
try:
|
||||
async with self.lock:
|
||||
if job["status"] == "cancelled":
|
||||
return
|
||||
self.queue_counts[job_type] = max(0, self.queue_counts[job_type] - 1)
|
||||
self.running_counts[job_type] += 1
|
||||
job["status"] = "running"
|
||||
job["updated_at"] = _now_ms()
|
||||
metrics = self._metrics(job_type)
|
||||
await self._publish(job_id, "started", {"job_id": job_id, "type": job_type, "status": "running", **metrics})
|
||||
|
||||
async def emit(event: str, data: dict[str, Any]) -> None:
|
||||
metrics_now = self._metrics(job_type)
|
||||
await self._publish(job_id, event, {"job_id": job_id, "type": job_type, "status": job["status"], **metrics_now, **data})
|
||||
|
||||
def is_cancelled() -> bool:
|
||||
return bool(job.get("cancel_requested"))
|
||||
|
||||
result = await self.handlers[job_type](job["payload"], emit, is_cancelled)
|
||||
async with self.lock:
|
||||
if job["cancel_requested"]:
|
||||
job["status"] = "cancelled"
|
||||
metrics = self._metrics(job_type)
|
||||
await emit("cancelled", {"reason": "abort"})
|
||||
return
|
||||
job["status"] = "completed"
|
||||
job["result"] = result
|
||||
job["updated_at"] = _now_ms()
|
||||
metrics = self._metrics(job_type)
|
||||
await self._publish(job_id, "done", {"job_id": job_id, "type": job_type, "status": "completed", "result": result, **metrics})
|
||||
except asyncio.CancelledError:
|
||||
async with self.lock:
|
||||
job["status"] = "cancelled"
|
||||
job["cancel_requested"] = True
|
||||
metrics = self._metrics(job_type)
|
||||
await self._publish(job_id, "cancelled", {"job_id": job_id, "type": job_type, "status": "cancelled", **metrics})
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.exception("in-memory job failed job_id=%s type=%s", job_id, job_type)
|
||||
async with self.lock:
|
||||
job["status"] = "failed"
|
||||
job["error"] = str(exc)
|
||||
job["updated_at"] = _now_ms()
|
||||
metrics = self._metrics(job_type)
|
||||
await self._publish(job_id, "error", {"job_id": job_id, "type": job_type, "status": "failed", "error": str(exc), **metrics})
|
||||
finally:
|
||||
async with self.lock:
|
||||
self.running_counts[job_type] = max(0, self.running_counts[job_type] - 1)
|
||||
self.running_tasks.pop(job_id, None)
|
||||
self.semaphores[job_type].release()
|
||||
|
||||
async def close(self) -> None:
|
||||
for task in self.worker_tasks:
|
||||
task.cancel()
|
||||
for task in self.running_tasks.values():
|
||||
task.cancel()
|
||||
for task in self.worker_tasks:
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
self.worker_tasks.clear()
|
||||
self.running_tasks.clear()
|
||||
self.started = False
|
||||
|
||||
|
||||
class RedisJobManager(BaseJobManager):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
if redis_asyncio is None:
|
||||
raise JobSystemError("redis package is not installed")
|
||||
self.redis = redis_asyncio.from_url(
|
||||
os.getenv("REDIS_URL", "redis://localhost:6379/0"),
|
||||
encoding="utf-8",
|
||||
decode_responses=True,
|
||||
)
|
||||
self.prefix = (os.getenv("JOB_REDIS_PREFIX") or "llmtext:jobs").strip() or "llmtext:jobs"
|
||||
self.state_ttl = _int_env("JOB_STATE_TTL_SECONDS", 600)
|
||||
self.event_ttl = _int_env("JOB_EVENT_TTL_SECONDS", 600)
|
||||
|
||||
def _queue_key(self, job_type: str) -> str:
|
||||
return f"{self.prefix}:queue:{job_type}"
|
||||
|
||||
def _event_key(self, job_id: str) -> str:
|
||||
return f"{self.prefix}:events:{job_id}"
|
||||
|
||||
def _state_key(self, job_id: str) -> str:
|
||||
return f"{self.prefix}:state:{job_id}"
|
||||
|
||||
def _metrics_key(self, job_type: str) -> str:
|
||||
return f"{self.prefix}:metrics:{job_type}"
|
||||
|
||||
def _group_name(self, job_type: str) -> str:
|
||||
return f"{self.prefix}:group:{job_type}"
|
||||
|
||||
async def ensure_groups(self) -> None:
|
||||
for job_type in JOB_TYPES:
|
||||
stream = self._queue_key(job_type)
|
||||
group = self._group_name(job_type)
|
||||
try:
|
||||
await self.redis.xgroup_create(stream, group, id="0-0", mkstream=True)
|
||||
except Exception as exc: # pragma: no cover - redis-specific
|
||||
if "BUSYGROUP" not in str(exc):
|
||||
raise
|
||||
|
||||
async def close(self) -> None:
|
||||
await self.redis.aclose()
|
||||
|
||||
async def _metrics(self, job_type: str) -> dict[str, Any]:
|
||||
raw = await self.redis.hgetall(self._metrics_key(job_type))
|
||||
queued = int(raw.get("queued_count", "0") or 0)
|
||||
running = int(raw.get("running_count", "0") or 0)
|
||||
config = _queue_config(job_type)
|
||||
capacity = max(config.max_queue + config.concurrency, 1)
|
||||
ratio = min((queued + running) / capacity, 1.0)
|
||||
return {
|
||||
"queued_count": queued,
|
||||
"running_count": running,
|
||||
"concurrency_limit": config.concurrency,
|
||||
"max_queue": config.max_queue,
|
||||
"busy_ratio": round(ratio, 4),
|
||||
"busy_level": _busy_level(ratio),
|
||||
}
|
||||
|
||||
async def _emit_event(self, job_id: str, event: str, data: dict[str, Any]) -> None:
|
||||
key = self._event_key(job_id)
|
||||
payload = {k: _json_dumps(v) if not isinstance(v, str) else v for k, v in data.items()}
|
||||
payload["event"] = event
|
||||
await self.redis.xadd(key, payload, maxlen=_int_env("JOB_EVENT_STREAM_MAXLEN", 512), approximate=True)
|
||||
await self.redis.expire(key, self.event_ttl)
|
||||
|
||||
async def _set_state(self, job_id: str, state: dict[str, Any]) -> None:
|
||||
serializable = {k: _json_dumps(v) if isinstance(v, (dict, list)) else str(v) for k, v in state.items()}
|
||||
await self.redis.hset(self._state_key(job_id), mapping=serializable)
|
||||
await self.redis.expire(self._state_key(job_id), self.state_ttl)
|
||||
|
||||
async def submit(self, job_type: str, payload: dict[str, Any], request_id: str | None = None) -> str:
|
||||
config = _queue_config(job_type)
|
||||
metrics = await self._metrics(job_type)
|
||||
if metrics["queued_count"] >= config.max_queue:
|
||||
raise QueueFullError(f"{job_type} queue is full")
|
||||
|
||||
job_id = request_id or str(uuid.uuid4())
|
||||
created_at = _now_ms()
|
||||
state = {
|
||||
"job_id": job_id,
|
||||
"request_id": job_id,
|
||||
"type": job_type,
|
||||
"status": "queued",
|
||||
"error": "",
|
||||
"created_at": created_at,
|
||||
"updated_at": created_at,
|
||||
"cancel_requested": "0",
|
||||
}
|
||||
await self._set_state(job_id, state)
|
||||
await self.redis.hincrby(self._metrics_key(job_type), "queued_count", 1)
|
||||
await self.redis.expire(self._metrics_key(job_type), self.state_ttl)
|
||||
metrics = await self._metrics(job_type)
|
||||
await self._emit_event(job_id, "queued", {"job_id": job_id, "type": job_type, "status": "queued", **metrics})
|
||||
await self.redis.xadd(self._queue_key(job_type), {"job_id": job_id, "payload": _json_dumps(payload), "request_id": job_id})
|
||||
return job_id
|
||||
|
||||
async def cancel(self, job_id: str, reason: str = "abort") -> dict[str, Any]:
|
||||
state = await self.get_status(job_id)
|
||||
if not state:
|
||||
return {"cancelled": False, "status": "not_found"}
|
||||
if state["status"] in TERMINAL_STATUSES:
|
||||
return {"cancelled": False, "status": state["status"]}
|
||||
await self.redis.hset(self._state_key(job_id), mapping={"cancel_requested": "1", "status": "cancelled", "updated_at": _now_ms(), "cancel_reason": reason})
|
||||
metrics = await self._metrics(state["type"])
|
||||
await self._emit_event(job_id, "cancelled", {"job_id": job_id, "type": state["type"], "status": "cancelled", "reason": reason, **metrics})
|
||||
return {"cancelled": True, "status": "ok"}
|
||||
|
||||
async def get_status(self, job_id: str) -> dict[str, Any] | None:
|
||||
state = await self.redis.hgetall(self._state_key(job_id))
|
||||
if not state:
|
||||
return None
|
||||
job_type = state.get("type", "")
|
||||
metrics = await self._metrics(job_type) if job_type else {}
|
||||
result = state.get("result")
|
||||
error = state.get("error", "")
|
||||
return {
|
||||
"job_id": state.get("job_id", job_id),
|
||||
"request_id": state.get("request_id", job_id),
|
||||
"type": job_type,
|
||||
"status": state.get("status", "queued"),
|
||||
"error": error,
|
||||
"result": _json_loads(result, result),
|
||||
"cancel_requested": state.get("cancel_requested") == "1",
|
||||
**metrics,
|
||||
}
|
||||
|
||||
async def stream_events(self, job_id: str) -> AsyncIterator[dict[str, Any]]:
|
||||
stream = self._event_key(job_id)
|
||||
last_id = "0-0"
|
||||
while True:
|
||||
events = await self.redis.xread({stream: last_id}, block=1000, count=20)
|
||||
if not events:
|
||||
state = await self.get_status(job_id)
|
||||
if state and state["status"] in TERMINAL_STATUSES:
|
||||
break
|
||||
continue
|
||||
for _, entries in events:
|
||||
for entry_id, fields in entries:
|
||||
last_id = entry_id
|
||||
event_payload: dict[str, Any] = {}
|
||||
for key, value in fields.items():
|
||||
if key == "event":
|
||||
event_payload[key] = value
|
||||
continue
|
||||
try:
|
||||
event_payload[key] = json.loads(value)
|
||||
except Exception:
|
||||
event_payload[key] = value
|
||||
yield event_payload
|
||||
if event_payload.get("event") in TERMINAL_EVENTS:
|
||||
return
|
||||
|
||||
|
||||
class RedisWorker:
|
||||
def __init__(self, manager: RedisJobManager) -> None:
|
||||
self.manager = manager
|
||||
self.running_tasks: dict[str, asyncio.Task] = {}
|
||||
self.queue_semaphores = {job_type: asyncio.Semaphore(_queue_config(job_type).concurrency) for job_type in JOB_TYPES}
|
||||
self.poll_interval = _float_env("JOB_CANCEL_POLL_SECONDS", 0.5)
|
||||
self.consumer_name = (os.getenv("JOB_CONSUMER_NAME") or f"worker-{uuid.uuid4().hex[:8]}").strip()
|
||||
|
||||
async def run_forever(self) -> None:
|
||||
await self.manager.ensure_groups()
|
||||
cancel_task = asyncio.create_task(self._cancel_watch_loop())
|
||||
consumers = [asyncio.create_task(self._consume_loop(job_type)) for job_type in JOB_TYPES]
|
||||
try:
|
||||
await asyncio.gather(*consumers)
|
||||
finally:
|
||||
cancel_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await cancel_task
|
||||
|
||||
async def _cancel_watch_loop(self) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
for job_id, task in list(self.running_tasks.items()):
|
||||
state = await self.manager.get_status(job_id)
|
||||
if state and state.get("cancel_requested") and not task.done():
|
||||
task.cancel()
|
||||
|
||||
async def _consume_loop(self, job_type: str) -> None:
|
||||
queue_key = self.manager._queue_key(job_type)
|
||||
group = self.manager._group_name(job_type)
|
||||
semaphore = self.queue_semaphores[job_type]
|
||||
while True:
|
||||
streams = await self.manager.redis.xreadgroup(group, self.consumer_name, {queue_key: ">"}, count=1, block=1000)
|
||||
if not streams:
|
||||
continue
|
||||
for _, messages in streams:
|
||||
for message_id, fields in messages:
|
||||
await semaphore.acquire()
|
||||
task = asyncio.create_task(self._run_message(job_type, queue_key, group, message_id, fields, semaphore))
|
||||
self.running_tasks[fields["job_id"]] = task
|
||||
|
||||
async def _run_message(
|
||||
self,
|
||||
job_type: str,
|
||||
queue_key: str,
|
||||
group: str,
|
||||
message_id: str,
|
||||
fields: dict[str, str],
|
||||
semaphore: asyncio.Semaphore,
|
||||
) -> None:
|
||||
job_id = fields["job_id"]
|
||||
try:
|
||||
state = await self.manager.get_status(job_id)
|
||||
if not state or state["status"] == "cancelled":
|
||||
await self.manager.redis.xack(queue_key, group, message_id)
|
||||
return
|
||||
|
||||
await self.manager.redis.hincrby(self.manager._metrics_key(job_type), "queued_count", -1)
|
||||
await self.manager.redis.hincrby(self.manager._metrics_key(job_type), "running_count", 1)
|
||||
await self.manager._set_state(job_id, {
|
||||
"job_id": job_id,
|
||||
"request_id": state["request_id"],
|
||||
"type": job_type,
|
||||
"status": "running",
|
||||
"updated_at": _now_ms(),
|
||||
"created_at": state.get("created_at", _now_ms()),
|
||||
"cancel_requested": "1" if state.get("cancel_requested") else "0",
|
||||
"error": "",
|
||||
})
|
||||
metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, "started", {"job_id": job_id, "type": job_type, "status": "running", **metrics})
|
||||
|
||||
payload = _json_loads(fields["payload"], {})
|
||||
|
||||
async def emit(event: str, data: dict[str, Any]) -> None:
|
||||
live_state = await self.manager.get_status(job_id) or {"status": "running"}
|
||||
live_metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, event, {"job_id": job_id, "type": job_type, "status": live_state["status"], **live_metrics, **data})
|
||||
|
||||
def is_cancelled() -> bool:
|
||||
task = self.running_tasks.get(job_id)
|
||||
return bool(task and task.cancelled())
|
||||
|
||||
result = await self.manager.handlers[job_type](payload, emit, is_cancelled)
|
||||
current = await self.manager.get_status(job_id)
|
||||
if current and current["status"] == "cancelled":
|
||||
return
|
||||
|
||||
await self.manager._set_state(job_id, {
|
||||
"job_id": job_id,
|
||||
"request_id": state["request_id"],
|
||||
"type": job_type,
|
||||
"status": "completed",
|
||||
"updated_at": _now_ms(),
|
||||
"created_at": state.get("created_at", _now_ms()),
|
||||
"cancel_requested": "0",
|
||||
"error": "",
|
||||
"result": _json_dumps(result),
|
||||
})
|
||||
metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, "done", {"job_id": job_id, "type": job_type, "status": "completed", "result": result, **metrics})
|
||||
await self.manager.redis.xack(queue_key, group, message_id)
|
||||
except asyncio.CancelledError:
|
||||
await self.manager.redis.hset(self.manager._state_key(job_id), mapping={"status": "cancelled", "cancel_requested": "1", "updated_at": _now_ms()})
|
||||
metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, "cancelled", {"job_id": job_id, "type": job_type, "status": "cancelled", **metrics})
|
||||
await self.manager.redis.xack(queue_key, group, message_id)
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.exception("redis worker failed job_id=%s type=%s", job_id, job_type)
|
||||
state = await self.manager.get_status(job_id)
|
||||
request_id = state["request_id"] if state else job_id
|
||||
await self.manager._set_state(job_id, {
|
||||
"job_id": job_id,
|
||||
"request_id": request_id,
|
||||
"type": job_type,
|
||||
"status": "failed",
|
||||
"updated_at": _now_ms(),
|
||||
"created_at": state.get("created_at", _now_ms()) if state else _now_ms(),
|
||||
"cancel_requested": "0",
|
||||
"error": str(exc),
|
||||
})
|
||||
metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, "error", {"job_id": job_id, "type": job_type, "status": "failed", "error": str(exc), **metrics})
|
||||
await self.manager.redis.xack(queue_key, group, message_id)
|
||||
finally:
|
||||
self.running_tasks.pop(job_id, None)
|
||||
await self.manager.redis.hincrby(self.manager._metrics_key(job_type), "running_count", -1)
|
||||
semaphore.release()
|
||||
|
||||
|
||||
_job_manager: BaseJobManager | None = None
|
||||
|
||||
|
||||
def get_job_manager() -> BaseJobManager:
|
||||
global _job_manager
|
||||
if _job_manager is None:
|
||||
backend = get_job_backend_name()
|
||||
if backend == "redis":
|
||||
_job_manager = RedisJobManager()
|
||||
else:
|
||||
_job_manager = InMemoryJobManager()
|
||||
return _job_manager
|
||||
|
||||
|
||||
def reset_job_manager() -> None:
|
||||
global _job_manager
|
||||
manager = _job_manager
|
||||
if manager is not None:
|
||||
close = getattr(manager, "close", None)
|
||||
if close is not None:
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
try:
|
||||
asyncio.run(close())
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
loop.create_task(close())
|
||||
_job_manager = None
|
||||
+10
-8
@@ -87,16 +87,17 @@ def _build_chat_payload(
|
||||
if prefill:
|
||||
messages.append({'role': 'assistant', 'content': prefill})
|
||||
|
||||
options = {'temperature': temperature}
|
||||
if thinking:
|
||||
options['think'] = thinking
|
||||
|
||||
payload = {
|
||||
'model': _resolve_model_name(model, use_pro_model=use_pro_model),
|
||||
'messages': messages,
|
||||
'stream': False,
|
||||
'options': options,
|
||||
}
|
||||
|
||||
options = {'temperature': temperature}
|
||||
if thinking:
|
||||
payload['options'] = {'temperature': temperature, 'think': thinking}
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
@@ -119,16 +120,17 @@ def _build_chat_stream_payload(
|
||||
if prefill:
|
||||
messages.append({'role': 'assistant', 'content': prefill})
|
||||
|
||||
options = {'temperature': temperature}
|
||||
if thinking:
|
||||
options['think'] = thinking
|
||||
|
||||
payload = {
|
||||
'model': _resolve_model_name(model, use_pro_model=use_pro_model),
|
||||
'messages': messages,
|
||||
'stream': True,
|
||||
'options': options,
|
||||
}
|
||||
|
||||
options = {'temperature': temperature}
|
||||
if thinking:
|
||||
payload['options'] = {'temperature': temperature, 'think': thinking}
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
|
||||
+326
-339
@@ -1,12 +1,7 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import uuid
|
||||
from typing import Optional
|
||||
|
||||
@@ -17,10 +12,28 @@ from fastapi.security import APIKeyHeader
|
||||
from pydantic import BaseModel
|
||||
|
||||
from geoip import get_ip_location_text
|
||||
from llm import call_ollama, call_vlm_ocr, stream_ollama
|
||||
from job_handlers import (
|
||||
_sanitize_converted_markdown,
|
||||
sanitize_inline_completion_content,
|
||||
ALLOWED_CONVERT_EXTENSIONS,
|
||||
asr_handler,
|
||||
completion_handler,
|
||||
compress_handler,
|
||||
convert_handler,
|
||||
ocr_handler,
|
||||
pro_completion_handler,
|
||||
tts_handler,
|
||||
)
|
||||
from job_system import (
|
||||
InMemoryJobManager,
|
||||
JobSystemError,
|
||||
JOB_TYPES,
|
||||
QueueFullError,
|
||||
RedisJobManager,
|
||||
get_job_manager,
|
||||
persist_temp_input,
|
||||
)
|
||||
from models import UserPreferences
|
||||
from prompt import build_completion_prompts, prepare_prompt_context
|
||||
import markitdown
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
@@ -28,22 +41,7 @@ logging.basicConfig(
|
||||
)
|
||||
logger = logging.getLogger("api")
|
||||
|
||||
_markitdown_instance = None
|
||||
|
||||
|
||||
def _get_markitdown(): # pragma: no cover
|
||||
global _markitdown_instance
|
||||
if _markitdown_instance is None:
|
||||
_markitdown_instance = markitdown.MarkItDown()
|
||||
return _markitdown_instance
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
# Startup event disabled — TTS model loads lazily on first request
|
||||
# to avoid blocking startup and OOM crashes.
|
||||
ACTIVE_COMPLETIONS: dict[str, asyncio.Task] = {}
|
||||
ACTIVE_COMPLETIONS_LOCK = asyncio.Lock()
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
@@ -53,16 +51,9 @@ app.add_middleware(
|
||||
)
|
||||
|
||||
API_KEY = os.getenv("API_KEY", "your-secret-key-here")
|
||||
DOC_COMPRESS_CONTEXT_LIMIT = int(os.getenv("DOC_COMPRESS_CONTEXT_LIMIT", "128000"))
|
||||
api_key_header = APIKeyHeader(name="X-API-Key")
|
||||
|
||||
|
||||
async def get_api_key(api_key: str = Security(api_key_header)): # pragma: no cover
|
||||
if api_key != API_KEY:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Could not validate credentials",
|
||||
)
|
||||
return api_key
|
||||
_handlers_registered = False
|
||||
|
||||
|
||||
class CompletionRequest(BaseModel):
|
||||
@@ -76,6 +67,16 @@ class CompletionRequest(BaseModel):
|
||||
temperature: float = 0.7
|
||||
|
||||
|
||||
class ProCompletionRequest(BaseModel):
|
||||
prefix: str
|
||||
suffix: str
|
||||
languageId: str = "markdown"
|
||||
instruction: str = ""
|
||||
pro_thinking: str = "medium"
|
||||
privacy_mode: bool = False
|
||||
user_preferences: Optional[UserPreferences] = None
|
||||
|
||||
|
||||
class CancelCompletionRequest(BaseModel):
|
||||
request_id: str
|
||||
reason: str = "abort"
|
||||
@@ -92,30 +93,21 @@ class ConvertRequest(BaseModel):
|
||||
filename: str = "document.pdf"
|
||||
|
||||
|
||||
ALLOWED_CONVERT_EXTENSIONS = {".txt", ".docx", ".pptx", ".pdf"}
|
||||
IMAGE_MARKDOWN_RE = re.compile(r"!\[[^\]]*]\([^)]+\)")
|
||||
IMAGE_HTML_RE = re.compile(r"<img\b[^>]*>", re.IGNORECASE)
|
||||
class CompressRequest(BaseModel):
|
||||
content: str
|
||||
docType: str = "txt"
|
||||
|
||||
|
||||
def _convert_docx_to_pdf(input_path: str, output_path: str) -> None: # pragma: no cover
|
||||
node_executable = shutil.which("node")
|
||||
if not node_executable:
|
||||
raise RuntimeError("未找到 Node.js,无法转换 DOCX 为 PDF")
|
||||
class TTSJobRequest(BaseModel):
|
||||
text: str
|
||||
instruct: str = ""
|
||||
speaker: str = "Vivian"
|
||||
format: str = "wav"
|
||||
|
||||
bridge_path = os.path.join(os.path.dirname(__file__), "docx2pdf_bridge.cjs")
|
||||
if not os.path.exists(bridge_path):
|
||||
raise RuntimeError("缺少 DOCX 转 PDF 桥接脚本")
|
||||
|
||||
result = subprocess.run(
|
||||
[node_executable, bridge_path, input_path, output_path],
|
||||
cwd=os.path.dirname(os.path.dirname(__file__)),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
error_text = (result.stderr or result.stdout or "DOCX 转 PDF 失败").strip()
|
||||
raise RuntimeError(error_text)
|
||||
class ASRJobRequest(BaseModel):
|
||||
audio_base64: str
|
||||
language: Optional[str] = "zh-CN"
|
||||
|
||||
|
||||
def _preview(text: str, limit: int = 80) -> str:
|
||||
@@ -125,14 +117,6 @@ def _preview(text: str, limit: int = 80) -> str:
|
||||
return value[:limit] + "..."
|
||||
|
||||
|
||||
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 get_client_ip(request: Request) -> str:
|
||||
if request.client:
|
||||
return request.headers.get("X-Client-IP") or request.client.host
|
||||
@@ -147,197 +131,56 @@ def _clamp_temperature(value: float, default: float = 0.7) -> float:
|
||||
return max(0.0, min(numeric, 1.2))
|
||||
|
||||
|
||||
@app.post("/v1/completions")
|
||||
async def create_completion(request: Request, req: CompletionRequest, api_key: str = Security(get_api_key)):
|
||||
request_id = request.headers.get("X-Request-Id") or str(uuid.uuid4())
|
||||
request_tag = request_id[:8]
|
||||
inference_task: Optional[asyncio.Task] = None
|
||||
|
||||
client_ip = "hidden"
|
||||
location = ""
|
||||
|
||||
if not req.privacy_mode: # pragma: no cover
|
||||
client_ip = get_client_ip(request)
|
||||
location = get_ip_location_text(client_ip)
|
||||
if location:
|
||||
logger.info("[%s] client_location=%s", request_tag, location)
|
||||
|
||||
try:
|
||||
logger.info(
|
||||
"[%s] /v1/completions request_id=%s client_ip=%s prefix_chars=%d suffix_chars=%d lang=%s thinking=%s privacy=%s",
|
||||
request_tag,
|
||||
request_id,
|
||||
client_ip,
|
||||
len(req.prefix or ""),
|
||||
len(req.suffix or ""),
|
||||
req.languageId,
|
||||
req.model_thinking,
|
||||
req.privacy_mode,
|
||||
)
|
||||
|
||||
llm_prefix, llm_suffix = prepare_prompt_context(req.prefix or "", req.suffix or "")
|
||||
logger.info("[%s] llm_input_prefix=%r", request_tag, llm_prefix)
|
||||
logger.info("[%s] llm_input_suffix=%r", request_tag, llm_suffix)
|
||||
|
||||
system_prompt, user_prompt, prefill = build_completion_prompts(
|
||||
req.prefix,
|
||||
req.suffix,
|
||||
req.languageId,
|
||||
location=location,
|
||||
thinking_level=req.model_thinking,
|
||||
preferences=req.user_preferences,
|
||||
)
|
||||
|
||||
inference_task = asyncio.create_task(
|
||||
call_ollama(
|
||||
user_prompt,
|
||||
system_prompt=system_prompt,
|
||||
tag=f"{request_tag}-primary",
|
||||
temperature=_clamp_temperature(req.temperature, 0.7),
|
||||
thinking=req.model_thinking if req.model_thinking != "none" else None,
|
||||
model=req.model,
|
||||
prefill=prefill or None,
|
||||
)
|
||||
)
|
||||
|
||||
existing = ACTIVE_COMPLETIONS.get(request_id)
|
||||
if existing and not existing.done():
|
||||
existing.cancel()
|
||||
ACTIVE_COMPLETIONS[request_id] = inference_task
|
||||
|
||||
result = await inference_task
|
||||
content = result["content"] or ""
|
||||
if not content.strip():
|
||||
logger.warning("[%s] primary returned empty content, returning empty result", request_tag)
|
||||
logger.info(
|
||||
"[%s] completion resolved source=primary request_id=%s content_chars=%d content_preview='%s'",
|
||||
request_tag,
|
||||
request_id,
|
||||
len(content),
|
||||
_preview(content, 120),
|
||||
)
|
||||
|
||||
return JSONResponse(content={"content": content, "request_id": request_id})
|
||||
except asyncio.CancelledError:
|
||||
logger.info("[%s] /v1/completions cancelled request_id=%s", request_tag, request_id)
|
||||
return JSONResponse(content={"cancelled": True, "request_id": request_id}, status_code=499)
|
||||
except Exception as e:
|
||||
logger.exception("[%s] /v1/completions failed request_id=%s: %s", request_tag, request_id, e)
|
||||
return JSONResponse(content={"error": str(e)}, status_code=500)
|
||||
finally:
|
||||
active = ACTIVE_COMPLETIONS.get(request_id)
|
||||
if active is not None and active is inference_task:
|
||||
ACTIVE_COMPLETIONS.pop(request_id, None)
|
||||
async def get_api_key(api_key: str = Security(api_key_header)): # pragma: no cover
|
||||
if api_key != API_KEY:
|
||||
raise HTTPException(status_code=403, detail="Could not validate credentials")
|
||||
return api_key
|
||||
|
||||
|
||||
@app.post("/v1/pro/completions/stream")
|
||||
async def create_pro_completion_stream(request: Request, req: CompletionRequest, api_key: str = Security(get_api_key)):
|
||||
request_id = request.headers.get("X-Request-Id") or str(uuid.uuid4())
|
||||
request_tag = request_id[:8]
|
||||
queue: asyncio.Queue[Optional[tuple[str, str]]] = asyncio.Queue()
|
||||
def _serialize_preferences(preferences: UserPreferences | None) -> dict | None:
|
||||
if preferences is None:
|
||||
return None
|
||||
if hasattr(preferences, "dict"):
|
||||
return preferences.dict()
|
||||
return dict(preferences)
|
||||
|
||||
client_ip = "hidden"
|
||||
location = ""
|
||||
|
||||
if not req.privacy_mode: # pragma: no cover
|
||||
client_ip = get_client_ip(request)
|
||||
location = get_ip_location_text(client_ip)
|
||||
if location:
|
||||
logger.info("[%s] client_location=%s", request_tag, location)
|
||||
def _request_id(request: Request) -> str:
|
||||
return request.headers.get("X-Request-Id") or str(uuid.uuid4())
|
||||
|
||||
logger.info(
|
||||
"[%s] /v1/pro/completions/stream request_id=%s client_ip=%s prefix_chars=%d suffix_chars=%d lang=%s thinking=%s privacy=%s model=%s temp=%.2f",
|
||||
request_tag,
|
||||
request_id,
|
||||
client_ip,
|
||||
len(req.prefix or ""),
|
||||
len(req.suffix or ""),
|
||||
req.languageId,
|
||||
req.model_thinking,
|
||||
req.privacy_mode,
|
||||
req.model or "",
|
||||
_clamp_temperature(req.temperature, 0.7),
|
||||
)
|
||||
|
||||
llm_prefix, llm_suffix = prepare_prompt_context(req.prefix or "", req.suffix or "")
|
||||
logger.info("[%s] pro_llm_input_prefix=%r", request_tag, llm_prefix)
|
||||
logger.info("[%s] pro_llm_input_suffix=%r", request_tag, llm_suffix)
|
||||
def _register_handlers() -> None:
|
||||
global _handlers_registered
|
||||
if _handlers_registered:
|
||||
return
|
||||
manager = get_job_manager()
|
||||
manager.register_handler("completion", completion_handler)
|
||||
manager.register_handler("pro_completion", pro_completion_handler)
|
||||
manager.register_handler("compress", compress_handler)
|
||||
manager.register_handler("ocr", ocr_handler)
|
||||
manager.register_handler("convert", convert_handler)
|
||||
manager.register_handler("tts", tts_handler)
|
||||
manager.register_handler("asr", asr_handler)
|
||||
_handlers_registered = True
|
||||
|
||||
system_prompt, user_prompt, prefill = build_completion_prompts(
|
||||
req.prefix,
|
||||
req.suffix,
|
||||
req.languageId,
|
||||
location=location,
|
||||
thinking_level=req.model_thinking,
|
||||
preferences=req.user_preferences,
|
||||
)
|
||||
|
||||
async def producer() -> None:
|
||||
chunks: list[str] = []
|
||||
try:
|
||||
async for delta in stream_ollama(
|
||||
user_prompt,
|
||||
system_prompt=system_prompt,
|
||||
tag=f"{request_tag}-pro",
|
||||
temperature=_clamp_temperature(req.temperature, 0.7),
|
||||
thinking=req.model_thinking if req.model_thinking != "none" else None,
|
||||
model=req.model,
|
||||
use_pro_model=True,
|
||||
prefill=prefill or None,
|
||||
):
|
||||
chunks.append(delta)
|
||||
await queue.put(("chunk", json.dumps({"delta": delta}, ensure_ascii=False)))
|
||||
def _sse(event: str, data: dict) -> str:
|
||||
return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
|
||||
|
||||
content = "".join(chunks)
|
||||
logger.info(
|
||||
"[%s] pro stream resolved request_id=%s content_chars=%d content_preview='%s'",
|
||||
request_tag,
|
||||
request_id,
|
||||
len(content),
|
||||
_preview(content, 120),
|
||||
)
|
||||
await queue.put((
|
||||
"done",
|
||||
json.dumps({"content": content, "request_id": request_id}, ensure_ascii=False),
|
||||
))
|
||||
except asyncio.CancelledError:
|
||||
logger.info("[%s] /v1/pro/completions/stream cancelled request_id=%s", request_tag, request_id)
|
||||
await queue.put((
|
||||
"cancelled",
|
||||
json.dumps({"cancelled": True, "request_id": request_id}, ensure_ascii=False),
|
||||
))
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.exception("[%s] /v1/pro/completions/stream failed request_id=%s: %s", request_tag, request_id, e)
|
||||
await queue.put((
|
||||
"error",
|
||||
json.dumps({"error": str(e), "request_id": request_id}, ensure_ascii=False),
|
||||
))
|
||||
finally:
|
||||
await queue.put(None)
|
||||
|
||||
producer_task = asyncio.create_task(producer())
|
||||
existing = ACTIVE_COMPLETIONS.get(request_id)
|
||||
if existing and not existing.done():
|
||||
existing.cancel()
|
||||
ACTIVE_COMPLETIONS[request_id] = producer_task
|
||||
async def _stream_job(job_id: str):
|
||||
_register_handlers()
|
||||
manager = get_job_manager()
|
||||
|
||||
async def event_stream():
|
||||
try:
|
||||
while True:
|
||||
item = await queue.get()
|
||||
if item is None:
|
||||
break
|
||||
|
||||
event_name, data = item
|
||||
yield f"event: {event_name}\ndata: {data}\n\n"
|
||||
except asyncio.CancelledError:
|
||||
producer_task.cancel()
|
||||
raise
|
||||
finally:
|
||||
active = ACTIVE_COMPLETIONS.get(request_id)
|
||||
if active is producer_task:
|
||||
ACTIVE_COMPLETIONS.pop(request_id, None)
|
||||
async for event in manager.stream_events(job_id):
|
||||
event_name = event.get("event", "message")
|
||||
payload = {k: v for k, v in event.items() if k != "event"}
|
||||
yield _sse(event_name, payload)
|
||||
except Exception as exc:
|
||||
logger.exception("job stream failed job_id=%s", job_id)
|
||||
yield _sse("error", {"job_id": job_id, "error": str(exc)})
|
||||
|
||||
return StreamingResponse(
|
||||
event_stream(),
|
||||
@@ -349,130 +192,266 @@ async def create_pro_completion_stream(request: Request, req: CompletionRequest,
|
||||
)
|
||||
|
||||
|
||||
async def _queue_job(job_type: str, payload: dict, request_id: str) -> str:
|
||||
_register_handlers()
|
||||
manager = get_job_manager()
|
||||
return await manager.submit(job_type, payload, request_id=request_id)
|
||||
|
||||
|
||||
async def _cancel_job(request_id: str, reason: str) -> dict:
|
||||
_register_handlers()
|
||||
manager = get_job_manager()
|
||||
return await manager.cancel(request_id, reason)
|
||||
|
||||
|
||||
async def _job_status(job_id: str) -> dict | None:
|
||||
_register_handlers()
|
||||
manager = get_job_manager()
|
||||
return await manager.get_status(job_id)
|
||||
|
||||
|
||||
async def _queue_load_snapshot() -> dict:
|
||||
manager = get_job_manager()
|
||||
if isinstance(manager, InMemoryJobManager):
|
||||
return {job_type: manager._metrics(job_type) for job_type in manager.queues}
|
||||
if isinstance(manager, RedisJobManager):
|
||||
return {job_type: await manager._metrics(job_type) for job_type in JOB_TYPES}
|
||||
return {}
|
||||
|
||||
|
||||
@app.post("/v1/completions")
|
||||
async def create_completion(
|
||||
request: Request,
|
||||
req: CompletionRequest,
|
||||
api_key: str = Security(get_api_key),
|
||||
):
|
||||
del api_key
|
||||
request_id = _request_id(request)
|
||||
location = ""
|
||||
if not req.privacy_mode: # pragma: no cover
|
||||
location = get_ip_location_text(get_client_ip(request))
|
||||
payload = {
|
||||
"request_id": request_id,
|
||||
"location": location,
|
||||
"request": {
|
||||
"prefix": req.prefix,
|
||||
"suffix": req.suffix,
|
||||
"languageId": req.languageId,
|
||||
"model_thinking": req.model_thinking,
|
||||
"privacy_mode": req.privacy_mode,
|
||||
"user_preferences": _serialize_preferences(req.user_preferences),
|
||||
"model": req.model,
|
||||
"temperature": _clamp_temperature(req.temperature, 0.7),
|
||||
},
|
||||
}
|
||||
try:
|
||||
job_id = await _queue_job("completion", payload, request_id)
|
||||
except QueueFullError as exc:
|
||||
return JSONResponse({"error": str(exc), "request_id": request_id}, status_code=429)
|
||||
except JobSystemError as exc:
|
||||
return JSONResponse({"error": str(exc), "request_id": request_id}, status_code=503)
|
||||
return await _stream_job(job_id)
|
||||
|
||||
|
||||
@app.post("/v1/completions/cancel")
|
||||
async def cancel_completion(req: CancelCompletionRequest, api_key: str = Security(get_api_key)):
|
||||
request_tag = str(uuid.uuid4())[:8]
|
||||
request_id = req.request_id or ""
|
||||
del api_key
|
||||
return await _cancel_job(req.request_id or "", req.reason)
|
||||
|
||||
async with ACTIVE_COMPLETIONS_LOCK:
|
||||
task = ACTIVE_COMPLETIONS.get(request_id)
|
||||
if task is None:
|
||||
logger.info(
|
||||
"[%s] /v1/completions/cancel request_id=%s status=not_found reason=%s",
|
||||
request_tag,
|
||||
request_id,
|
||||
req.reason,
|
||||
)
|
||||
return {"cancelled": False, "status": "not_found"}
|
||||
|
||||
if task.done():
|
||||
logger.info(
|
||||
"[%s] /v1/completions/cancel request_id=%s status=already_done reason=%s",
|
||||
request_tag,
|
||||
request_id,
|
||||
req.reason,
|
||||
)
|
||||
return {"cancelled": False, "status": "already_done"}
|
||||
@app.post("/v1/pro/completions")
|
||||
async def create_pro_completion(
|
||||
request: Request,
|
||||
req: ProCompletionRequest,
|
||||
api_key: str = Security(get_api_key),
|
||||
):
|
||||
del api_key
|
||||
request_id = _request_id(request)
|
||||
location = ""
|
||||
if not req.privacy_mode: # pragma: no cover
|
||||
location = get_ip_location_text(get_client_ip(request))
|
||||
payload = {
|
||||
"request_id": request_id,
|
||||
"location": location,
|
||||
"request": {
|
||||
"prefix": req.prefix,
|
||||
"suffix": req.suffix,
|
||||
"languageId": req.languageId,
|
||||
"instruction": req.instruction,
|
||||
"pro_thinking": req.pro_thinking,
|
||||
"privacy_mode": req.privacy_mode,
|
||||
"user_preferences": _serialize_preferences(req.user_preferences),
|
||||
},
|
||||
}
|
||||
try:
|
||||
job_id = await _queue_job("pro_completion", payload, request_id)
|
||||
except QueueFullError as exc:
|
||||
return JSONResponse({"error": str(exc), "request_id": request_id}, status_code=429)
|
||||
except JobSystemError as exc:
|
||||
return JSONResponse({"error": str(exc), "request_id": request_id}, status_code=503)
|
||||
return await _stream_job(job_id)
|
||||
|
||||
task.cancel()
|
||||
|
||||
logger.info(
|
||||
"[%s] /v1/completions/cancel request_id=%s status=ok reason=%s",
|
||||
request_tag,
|
||||
request_id,
|
||||
req.reason,
|
||||
)
|
||||
return {"cancelled": True, "status": "ok"}
|
||||
@app.post("/v1/pro/completions/cancel")
|
||||
async def cancel_pro_completion(req: CancelCompletionRequest, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
return await _cancel_job(req.request_id or "", req.reason)
|
||||
|
||||
|
||||
@app.get("/v1/pro/completions/status/{request_id}")
|
||||
async def get_pro_completion_status(request_id: str, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
state = await _job_status(request_id)
|
||||
if state is None:
|
||||
raise HTTPException(status_code=404, detail="PRO request not found")
|
||||
return state
|
||||
|
||||
|
||||
@app.post("/v1/ocr")
|
||||
async def ocr_image(request: OCRRequest, api_key: str = Security(get_api_key)):
|
||||
request_id = str(uuid.uuid4())[:8]
|
||||
async def ocr_image(req: OCRRequest, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
request_id = str(uuid.uuid4())
|
||||
try:
|
||||
logger.info(
|
||||
"[%s] /v1/ocr filename=%s language=%s image_base64_chars=%d",
|
||||
request_id,
|
||||
request.filename,
|
||||
request.language,
|
||||
len(request.image or ""),
|
||||
)
|
||||
image_bytes = base64.b64decode(request.image)
|
||||
logger.info("[%s] /v1/ocr decoded image_bytes=%d", request_id, len(image_bytes))
|
||||
result = await call_vlm_ocr(image_bytes, request.language)
|
||||
logger.info(
|
||||
"[%s] /v1/ocr success text_chars=%d text_preview='%s'",
|
||||
request_id,
|
||||
len(result or ""),
|
||||
_preview(result or "", 120),
|
||||
)
|
||||
return {"text": result, "filename": request.filename}
|
||||
except Exception as e:
|
||||
logger.exception("[%s] /v1/ocr failed: %s", request_id, e)
|
||||
return JSONResponse(content={"error": str(e)}, status_code=500)
|
||||
image_bytes = base64.b64decode(req.image)
|
||||
except Exception as exc:
|
||||
return JSONResponse({"error": str(exc)}, status_code=500)
|
||||
input_path = persist_temp_input(image_bytes, os.path.splitext(req.filename)[1] or ".img")
|
||||
try:
|
||||
job_id = await _queue_job("ocr", {
|
||||
"request_id": request_id,
|
||||
"input_path": input_path,
|
||||
"filename": req.filename,
|
||||
"language": req.language,
|
||||
}, request_id)
|
||||
except Exception:
|
||||
if os.path.exists(input_path):
|
||||
os.unlink(input_path)
|
||||
raise
|
||||
return await _stream_job(job_id)
|
||||
|
||||
|
||||
@app.post("/v1/convert")
|
||||
async def convert_to_markdown(request: ConvertRequest, api_key: str = Security(get_api_key)):
|
||||
"""Convert file to markdown"""
|
||||
request_id = str(uuid.uuid4())[:8]
|
||||
|
||||
async def convert_to_markdown(req: ConvertRequest, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
request_id = str(uuid.uuid4())
|
||||
ext = os.path.splitext(req.filename)[1].lower()
|
||||
if ext not in ALLOWED_CONVERT_EXTENSIONS:
|
||||
return JSONResponse({"error": "仅支持 txt、docx、pptx、pdf 格式"}, status_code=500)
|
||||
try:
|
||||
logger.info(
|
||||
"[%s] /v1/convert filename=%s file_base64_chars=%d",
|
||||
request_id,
|
||||
request.filename,
|
||||
len(request.file or ""),
|
||||
file_bytes = base64.b64decode(req.file)
|
||||
except Exception as exc:
|
||||
return JSONResponse({"error": str(exc)}, status_code=500)
|
||||
input_path = persist_temp_input(file_bytes, ext or ".bin")
|
||||
try:
|
||||
job_id = await _queue_job("convert", {
|
||||
"request_id": request_id,
|
||||
"input_path": input_path,
|
||||
"filename": req.filename,
|
||||
}, request_id)
|
||||
except Exception:
|
||||
if os.path.exists(input_path):
|
||||
os.unlink(input_path)
|
||||
raise
|
||||
return await _stream_job(job_id)
|
||||
|
||||
|
||||
@app.post("/v1/compress/submit")
|
||||
async def submit_compress(req: CompressRequest, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
content = req.content or ""
|
||||
if not content.strip():
|
||||
raise HTTPException(status_code=400, detail="文档内容为空,无法压缩")
|
||||
if len(content) > DOC_COMPRESS_CONTEXT_LIMIT:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"文档内容过长({len(content)} 字符),超过限制 {DOC_COMPRESS_CONTEXT_LIMIT},无法压缩",
|
||||
)
|
||||
task_id = str(uuid.uuid4())
|
||||
await _queue_job("compress", {"request_id": task_id, "content": content, "docType": req.docType or "txt"}, task_id)
|
||||
return {"task_id": task_id, "status": "queued"}
|
||||
|
||||
# Decode base64
|
||||
file_bytes = base64.b64decode(request.file)
|
||||
logger.info("[%s] /v1/convert decoded file_bytes=%d", request_id, len(file_bytes))
|
||||
|
||||
# Get file extension
|
||||
ext = os.path.splitext(request.filename)[1].lower()
|
||||
@app.get("/v1/compress/status")
|
||||
async def get_compress_status(task_id: str, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
if not task_id:
|
||||
raise HTTPException(status_code=400, detail="缺少 task_id 参数")
|
||||
state = await _job_status(task_id)
|
||||
if state is None:
|
||||
raise HTTPException(status_code=404, detail="任务不存在或已过期")
|
||||
if state["status"] == "completed":
|
||||
result = state.get("result") or {}
|
||||
return {"task_id": task_id, "status": "completed", "content": result.get("content", "")}
|
||||
if state["status"] == "failed":
|
||||
return {"task_id": task_id, "status": "error", "message": state.get("error") or ""}
|
||||
if state["status"] == "cancelled":
|
||||
return {"task_id": task_id, "status": "cancelled"}
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": "processing" if state["status"] == "running" else "queued",
|
||||
"busy_level": state.get("busy_level"),
|
||||
"queued_count": state.get("queued_count"),
|
||||
"running_count": state.get("running_count"),
|
||||
}
|
||||
|
||||
if ext not in ALLOWED_CONVERT_EXTENSIONS:
|
||||
raise ValueError("仅支持 txt、docx、pptx、pdf 格式")
|
||||
|
||||
if ext == ".txt":
|
||||
markdown_text = _sanitize_converted_markdown(file_bytes.decode("utf-8", errors="ignore"))
|
||||
return {
|
||||
"markdown": markdown_text,
|
||||
"filename": request.filename
|
||||
}
|
||||
@app.post("/v1/tts-asr/tts")
|
||||
async def queue_tts(req: TTSJobRequest, request: Request, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
request_id = _request_id(request)
|
||||
job_id = await _queue_job("tts", {
|
||||
"request_id": request_id,
|
||||
"text": req.text,
|
||||
"instruct": req.instruct,
|
||||
"speaker": req.speaker,
|
||||
"format": req.format,
|
||||
}, request_id)
|
||||
return await _stream_job(job_id)
|
||||
|
||||
# Create temporary file
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
|
||||
tmp.write(file_bytes)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
# Convert using MarkItDown
|
||||
md = _get_markitdown()
|
||||
result = await asyncio.to_thread(md.convert, tmp_path)
|
||||
markdown_text = _sanitize_converted_markdown(result.text_content)
|
||||
@app.post("/v1/tts-asr/asr")
|
||||
async def queue_asr(req: ASRJobRequest, request: Request, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
request_id = _request_id(request)
|
||||
try:
|
||||
audio_bytes = base64.b64decode(req.audio_base64)
|
||||
except Exception as exc:
|
||||
return JSONResponse({"error": str(exc)}, status_code=500)
|
||||
input_path = persist_temp_input(audio_bytes, ".wav")
|
||||
try:
|
||||
job_id = await _queue_job("asr", {
|
||||
"request_id": request_id,
|
||||
"input_path": input_path,
|
||||
"language": req.language or "zh-CN",
|
||||
}, request_id)
|
||||
except Exception:
|
||||
if os.path.exists(input_path):
|
||||
os.unlink(input_path)
|
||||
raise
|
||||
return await _stream_job(job_id)
|
||||
|
||||
logger.info(
|
||||
"[%s] /v1/convert success text_chars=%d text_preview='%s'",
|
||||
request_id,
|
||||
len(markdown_text or ""),
|
||||
_preview(markdown_text, 120),
|
||||
)
|
||||
|
||||
return {
|
||||
"markdown": markdown_text,
|
||||
"filename": request.filename
|
||||
}
|
||||
finally:
|
||||
# Clean up temporary file
|
||||
if os.path.exists(tmp_path):
|
||||
os.unlink(tmp_path)
|
||||
@app.post("/v1/jobs/{job_id}/cancel")
|
||||
async def cancel_job(job_id: str, req: CancelCompletionRequest, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
return await _cancel_job(req.request_id or job_id, req.reason)
|
||||
|
||||
|
||||
@app.get("/v1/jobs/{job_id}/status")
|
||||
async def get_job_status(job_id: str, api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
state = await _job_status(job_id)
|
||||
if state is None:
|
||||
raise HTTPException(status_code=404, detail="job not found")
|
||||
return state
|
||||
|
||||
|
||||
@app.get("/v1/jobs/load")
|
||||
async def get_job_load(api_key: str = Security(get_api_key)):
|
||||
del api_key
|
||||
return {"queues": await _queue_load_snapshot()}
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("[%s] /v1/convert failed: %s", request_id, e)
|
||||
return JSONResponse(content={"error": str(e)}, status_code=500)
|
||||
|
||||
# TTS and ASR routes (lazy loaded to avoid heavy import on startup)
|
||||
def _register_tts_asr_routes():
|
||||
try:
|
||||
from tts_asr import register_tts_asr_routes
|
||||
@@ -484,14 +463,22 @@ def _register_tts_asr_routes():
|
||||
return
|
||||
|
||||
try:
|
||||
register_tts_asr_routes(app)
|
||||
register_tts_asr_routes(app, include_generation_routes=False)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to register TTS/ASR routes: %s", exc)
|
||||
|
||||
|
||||
_register_tts_asr_routes()
|
||||
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def _shutdown_job_manager(): # pragma: no cover
|
||||
manager = get_job_manager()
|
||||
close = getattr(manager, "close", None)
|
||||
if close is not None:
|
||||
await close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host="0.0.0.0", port=8001)
|
||||
|
||||
|
||||
@@ -260,15 +260,18 @@ def register_pro_completion_routes(app: FastAPI, get_api_key):
|
||||
enable_thinking=True,
|
||||
timeout=PRO_COMPLETION_TIMEOUT,
|
||||
):
|
||||
# Handle 'thinking' event - just update state, don't accumulate
|
||||
if event_type == "thinking":
|
||||
await _send_sse_event(event_queue, "thinking", {"request_id": request_id})
|
||||
continue
|
||||
|
||||
# Handle 'chunk' event - accumulate content
|
||||
if not payload:
|
||||
continue
|
||||
chunks.append(payload)
|
||||
await _send_sse_event(event_queue, "chunk", {"delta": payload, "request_id": request_id})
|
||||
|
||||
# Handle 'done' event - return full content
|
||||
content = "".join(chunks)
|
||||
async with PRO_STATES_LOCK:
|
||||
if state.cancel_requested:
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
{
|
||||
"template": "You are the [PRO] model for LLM-IN-TEXT, specializing in high-precision markdown insertion for a {language_id} editor.\n\nReturn only the insertion text that should be placed between PREFIX and SUFFIX.\n\nPRO CORE PRINCIPLE:\n- Output insertion text only. No explanations, no analysis, no labels, no wrapper quotes.\n- Never output chain-of-thought or internal reasoning.\n- Never output control markers like <|fim_prefix|>, <|fim_suffix|>, <|fim_middle|>, assistant, final, channel.\n\nPRO MODE INTENT:\n- This is PRO_MODE=true. You may produce longer, structured markdown when instruction requires it.\n- Prioritize instruction fidelity first, then boundary safety, then style continuity.\n- If instruction is vague, continue naturally with concrete and useful content.\n\nBOUNDARY AND CONTEXT RULES:\n- Respect CURSOR_IN_FENCED_CODE_BLOCK, CURSOR_FENCE_LANGUAGE, MERMAID_CONTEXT, PREFIX_ENDS_WITH_NEWLINE, and SUFFIX_STARTS_WITH_NEWLINE.\n- Never repeat text from the beginning of SUFFIX.\n- Use minimum necessary newlines to avoid boundary collision.\n- Match PREFIX tone, language, and formatting conventions.\n\nSYNTAX PRIORITY:\n- Code block contexts must keep valid syntax and indentation.\n- Math must use $...$ for inline and $$...$$ for blocks unless inside latex fences.\n- Mermaid contexts must output valid mermaid statements; do not duplicate fences when already inside one.\n\nHIDDEN CONTEXT SAFETY:\n- OCR metadata and document-side context are hidden hints only.\n- Never copy hidden tags (e.g., <OCR:...>) into output.\n\nQUALITY BAR FOR PRO:\n- Prefer specific, information-dense output over generic filler.\n- For structured requests, preserve headings/list hierarchy and produce coherent section flow.\n- Keep output directly insertable without post-edit cleanups."
|
||||
"template": "You are the [PRO] model for LLM-IN-TEXT, specializing in high-precision markdown insertion for a {language_id} editor.\n\nReturn only the insertion text that should be placed between PREFIX and SUFFIX.\n\nPRO CORE PRINCIPLE:\n- Output insertion text only. No explanations, no analysis, no labels, no wrapper quotes.\n- Never wrap the entire answer in an outer ```markdown code fence; only use fenced code blocks when the inserted content itself requires code.\n- Never output chain-of-thought or internal reasoning.\n- Never output control markers like <|fim_prefix|>, <|fim_suffix|>, <|fim_middle|>, assistant, final, channel.\n\nPRO MODE INTENT:\n- This is PRO_MODE=true. You may produce longer, structured markdown when instruction requires it.\n- Prioritize instruction fidelity first, then boundary safety, then style continuity.\n- If instruction is vague, continue naturally with concrete and useful content.\n\nBOUNDARY AND CONTEXT RULES:\n- Respect CURSOR_IN_FENCED_CODE_BLOCK, CURSOR_FENCE_LANGUAGE, MERMAID_CONTEXT, PREFIX_ENDS_WITH_NEWLINE, and SUFFIX_STARTS_WITH_NEWLINE.\n- Never repeat text from the beginning of SUFFIX.\n- Use minimum necessary newlines to avoid boundary collision.\n- Match PREFIX tone, language, and formatting conventions.\n\nSYNTAX PRIORITY:\n- Code block contexts must keep valid syntax and indentation.\n- Math must use $...$ for inline and $$...$$ for blocks unless inside latex fences.\n- Mermaid contexts must output valid mermaid statements; do not duplicate fences when already inside one.\n\nHIDDEN CONTEXT SAFETY:\n- OCR metadata and document-side context are hidden hints only.\n- Never copy hidden tags (e.g., <OCR:...>) into output.\n\nQUALITY BAR FOR PRO:\n- Prefer specific, information-dense output over generic filler.\n- For structured requests, preserve headings/list hierarchy and produce coherent section flow.\n- Keep output directly insertable without post-edit cleanups."
|
||||
}
|
||||
|
||||
@@ -2,6 +2,8 @@ fastapi>=0.95.0
|
||||
uvicorn[standard]>=0.23.0
|
||||
pydantic>=1.10.0
|
||||
httpx>=0.24.0
|
||||
redis>=5.0.0
|
||||
python-dotenv>=1.0.0
|
||||
|
||||
numpy>=1.23.0
|
||||
soundfile>=0.10.3
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
os.environ["JOB_BACKEND"] = "memory"
|
||||
|
||||
CURRENT_DIR = Path(__file__).resolve().parent
|
||||
BACKEND_DIR = CURRENT_DIR.parent
|
||||
if str(BACKEND_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(BACKEND_DIR))
|
||||
|
||||
import job_handlers # type: ignore
|
||||
import job_system # type: ignore
|
||||
|
||||
main = importlib.import_module("main")
|
||||
|
||||
API_KEY = main.API_KEY
|
||||
HEADERS = {"X-API-Key": API_KEY}
|
||||
|
||||
|
||||
def setup_function():
|
||||
job_system.reset_job_manager()
|
||||
main._handlers_registered = False
|
||||
|
||||
|
||||
def _submit(client, content="test document", doc_type="txt", headers=None):
|
||||
return client.post("/v1/compress/submit", headers=headers if headers is not None else HEADERS, json={
|
||||
"content": content,
|
||||
"docType": doc_type,
|
||||
})
|
||||
|
||||
|
||||
def _status(client, task_id, headers=None):
|
||||
return client.get(f"/v1/compress/status?task_id={task_id}", headers=headers if headers is not None else HEADERS)
|
||||
|
||||
|
||||
def test_submit_empty_content_returns_400():
|
||||
with TestClient(main.app) as client:
|
||||
resp = _submit(client, "")
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_submit_too_long_returns_400(monkeypatch):
|
||||
monkeypatch.setattr(main, "DOC_COMPRESS_CONTEXT_LIMIT", 10)
|
||||
with TestClient(main.app) as client:
|
||||
resp = _submit(client, "a" * 100)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_submit_success_returns_task_id():
|
||||
with TestClient(main.app) as client:
|
||||
resp = _submit(client, "hello world")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert "task_id" in data
|
||||
assert data["status"] == "queued"
|
||||
|
||||
|
||||
def test_status_not_found_returns_404():
|
||||
with TestClient(main.app) as client:
|
||||
resp = _status(client, "nonexistent-id")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
def test_status_completed(monkeypatch):
|
||||
async def fake_call_ollama(prompt, system_prompt=None, **kwargs): # noqa: ARG001
|
||||
return {"content": f"[compressed] {prompt[:20]}"}
|
||||
|
||||
monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama)
|
||||
with TestClient(main.app) as client:
|
||||
resp = _submit(client, "important document text")
|
||||
task_id = resp.json()["task_id"]
|
||||
status_resp = _status(client, task_id)
|
||||
data = status_resp.json()
|
||||
assert status_resp.status_code == 200
|
||||
assert data["status"] in {"queued", "processing", "completed"}
|
||||
@@ -1,26 +1,31 @@
|
||||
import asyncio
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
os.environ["JOB_BACKEND"] = "memory"
|
||||
|
||||
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(BACKEND_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(BACKEND_DIR))
|
||||
|
||||
try:
|
||||
main = importlib.import_module("main")
|
||||
except ModuleNotFoundError:
|
||||
pytest.skip("main module dependencies are not available", allow_module_level=True)
|
||||
import job_handlers # type: ignore
|
||||
import job_system # type: ignore
|
||||
|
||||
main = importlib.import_module("main")
|
||||
|
||||
API_KEY_HEADERS = {"X-API-Key": "your-secret-key-here"}
|
||||
|
||||
|
||||
def setup_function():
|
||||
job_system.reset_job_manager()
|
||||
main._handlers_registered = False
|
||||
|
||||
|
||||
def _completion_payload():
|
||||
return {
|
||||
"prefix": "hello",
|
||||
@@ -32,7 +37,6 @@ def _completion_payload():
|
||||
|
||||
|
||||
def test_cancel_endpoint_cancels_running_task(monkeypatch):
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
started = threading.Event()
|
||||
cancelled = threading.Event()
|
||||
|
||||
@@ -45,21 +49,21 @@ def test_cancel_endpoint_cancels_running_task(monkeypatch):
|
||||
cancelled.set()
|
||||
raise
|
||||
|
||||
monkeypatch.setattr(main, "call_ollama", fake_call_ollama)
|
||||
monkeypatch.setattr(main, "build_completion_prompts", lambda *a, **k: ("system", "user"))
|
||||
monkeypatch.setattr(main, "prepare_prompt_context", lambda *a, **k: ("prefix", "suffix"))
|
||||
monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama)
|
||||
request_id = "req-cancel-1"
|
||||
|
||||
with TestClient(main.app) as client:
|
||||
request_id = "req-cancel-1"
|
||||
completion_headers = {**API_KEY_HEADERS, "X-Request-Id": request_id}
|
||||
response_box = {}
|
||||
|
||||
def send_completion():
|
||||
response_box["response"] = client.post(
|
||||
with client.stream(
|
||||
"POST",
|
||||
"/v1/completions",
|
||||
headers=completion_headers,
|
||||
headers={**API_KEY_HEADERS, "X-Request-Id": request_id},
|
||||
json=_completion_payload(),
|
||||
)
|
||||
) as response:
|
||||
response_box["status_code"] = response.status_code
|
||||
response_box["body"] = "".join(response.iter_text())
|
||||
|
||||
completion_thread = threading.Thread(target=send_completion, daemon=True)
|
||||
completion_thread.start()
|
||||
@@ -77,16 +81,10 @@ def test_cancel_endpoint_cancels_running_task(monkeypatch):
|
||||
completion_thread.join(timeout=5.0)
|
||||
assert not completion_thread.is_alive()
|
||||
assert cancelled.wait(timeout=2.0)
|
||||
|
||||
completion_response = response_box["response"]
|
||||
# 499 = client disconnected (TestClient timeout during cancel)
|
||||
assert completion_response.status_code in (200, 499)
|
||||
if completion_response.status_code == 200:
|
||||
assert completion_response.json()["cancelled"] is True
|
||||
assert "event: cancelled" in response_box["body"]
|
||||
|
||||
|
||||
def test_cancel_not_found():
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
with TestClient(main.app) as client:
|
||||
response = client.post(
|
||||
"/v1/completions/cancel",
|
||||
@@ -95,27 +93,3 @@ def test_cancel_not_found():
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"cancelled": False, "status": "not_found"}
|
||||
|
||||
|
||||
def test_completion_normal_flow(monkeypatch):
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
|
||||
async def fake_call_ollama(*args, **kwargs):
|
||||
return {"content": "completion text", "think": ""}
|
||||
|
||||
monkeypatch.setattr(main, "call_ollama", fake_call_ollama)
|
||||
monkeypatch.setattr(main, "build_completion_prompts", lambda *a, **k: ("system", "user"))
|
||||
monkeypatch.setattr(main, "prepare_prompt_context", lambda *a, **k: ("prefix", "suffix"))
|
||||
|
||||
with TestClient(main.app) as client:
|
||||
response = client.post(
|
||||
"/v1/completions",
|
||||
headers=API_KEY_HEADERS,
|
||||
json=_completion_payload(),
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["content"] == "completion text"
|
||||
assert data["request_id"] is not None
|
||||
assert main.ACTIVE_COMPLETIONS == {}
|
||||
|
||||
@@ -1,35 +1,29 @@
|
||||
import base64
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import base64
|
||||
import types
|
||||
import pytest
|
||||
from unittest.mock import MagicMock
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
BACKEND_DIR = os.path.abspath(os.path.join(CURRENT_DIR, ".."))
|
||||
if BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, BACKEND_DIR)
|
||||
os.environ["JOB_BACKEND"] = "memory"
|
||||
|
||||
if "tts_asr" not in sys.modules:
|
||||
fake_tts_asr = types.ModuleType("tts_asr")
|
||||
fake_tts_asr.register_tts_asr_routes = lambda app: None
|
||||
sys.modules["tts_asr"] = fake_tts_asr
|
||||
CURRENT_DIR = Path(__file__).resolve().parent
|
||||
BACKEND_DIR = CURRENT_DIR.parent
|
||||
if str(BACKEND_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(BACKEND_DIR))
|
||||
|
||||
import main # type: ignore
|
||||
import pro_completions # type: ignore
|
||||
import job_handlers # type: ignore
|
||||
import job_system # type: ignore
|
||||
|
||||
API_KEY = main.API_KEY
|
||||
HEADERS = {"X-API-Key": API_KEY}
|
||||
main = importlib.import_module("main")
|
||||
|
||||
HEADERS = {"X-API-Key": main.API_KEY}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_active_completions():
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
pro_completions.PRO_STATES.clear()
|
||||
yield
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
pro_completions.PRO_STATES.clear()
|
||||
def setup_function():
|
||||
job_system.reset_job_manager()
|
||||
main._handlers_registered = False
|
||||
|
||||
|
||||
class DummyRequest:
|
||||
@@ -51,64 +45,12 @@ def test_preview_long_text_truncated():
|
||||
assert main._preview(long_text) == long_text[:80] + "..."
|
||||
|
||||
|
||||
def test_preview_none_input():
|
||||
assert main._preview(None) == ""
|
||||
|
||||
|
||||
def test_preview_newlines_replaced():
|
||||
assert main._preview("line1\nline2") == "line1\\nline2"
|
||||
|
||||
|
||||
def test_sanitize_markdown_strips_image_markdown():
|
||||
assert "" not in main._sanitize_converted_markdown(
|
||||
"text with image  end"
|
||||
)
|
||||
|
||||
|
||||
def test_sanitize_markdown_strips_img_tag():
|
||||
assert "<img" not in main._sanitize_converted_markdown("<img src='x.png'/>")
|
||||
|
||||
|
||||
def test_sanitize_markdown_collapse_newlines():
|
||||
assert main._sanitize_converted_markdown("a\n\n\nb\n\n\n\nc") == "a\n\nb\n\nc"
|
||||
|
||||
|
||||
def test_sanitize_markdown_normalize_crlf():
|
||||
result = main._sanitize_converted_markdown("line1\r\nline2\r\n")
|
||||
assert "line1\nline2" in result
|
||||
assert "\r" not in result
|
||||
assert "" not in main._sanitize_converted_markdown("text ")
|
||||
|
||||
|
||||
def test_sanitize_inline_completion_strips_prefill():
|
||||
assert main.sanitize_inline_completion_content(
|
||||
"系统非常适合写作",
|
||||
prefill="系统",
|
||||
) == "非常适合写作"
|
||||
|
||||
|
||||
def test_sanitize_inline_completion_extracts_fim_middle():
|
||||
assert main.sanitize_inline_completion_content(
|
||||
"<|fim_middle|>系统非常适合写作<|end|>",
|
||||
prefill="系统",
|
||||
) == "非常适合写作"
|
||||
|
||||
|
||||
def test_sanitize_inline_completion_extracts_polluted_chat_output():
|
||||
polluted = (
|
||||
"on new line? Prefix ends with newline already. The suffix starts with no newline. "
|
||||
"We need to consider if output should end with newline? The suffix starts with no newline. "
|
||||
"So we output: \"让我们一起探索 AI 的无限可能。\""
|
||||
"<|end|><|start|>assistant<|channel|>final|fim_middle|>系统让我们一起探索 AI 的无限可能。"
|
||||
)
|
||||
assert main.sanitize_inline_completion_content(
|
||||
polluted,
|
||||
prefill="系统",
|
||||
) == "让我们一起探索 AI 的无限可能。"
|
||||
|
||||
|
||||
def test_get_client_ip_from_host():
|
||||
req = DummyRequest(host="1.2.3.4", headers={})
|
||||
assert main.get_client_ip(req) == "1.2.3.4"
|
||||
assert main.sanitize_inline_completion_content("系统非常适合写作", prefill="系统") == "非常适合写作"
|
||||
|
||||
|
||||
def test_get_client_ip_header_overrides_host():
|
||||
@@ -116,191 +58,64 @@ def test_get_client_ip_header_overrides_host():
|
||||
assert main.get_client_ip(req) == "5.6.7.8"
|
||||
|
||||
|
||||
def test_get_client_ip_when_client_missing():
|
||||
req = DummyRequest(host=None, headers={"X-Client-IP": "9.9.9.9"})
|
||||
req.client = None
|
||||
assert main.get_client_ip(req) == "9.9.9.9"
|
||||
|
||||
|
||||
def test_post_completions_wrong_api_key_returns_401():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions", json={
|
||||
"prefix": "hello", "suffix": "", "languageId": "markdown",
|
||||
"model_thinking": "low", "privacy_mode": True,
|
||||
})
|
||||
with TestClient(main.app) as client:
|
||||
resp = client.post("/v1/completions", json={
|
||||
"prefix": "hello", "suffix": "", "languageId": "markdown",
|
||||
"model_thinking": "low", "privacy_mode": True,
|
||||
})
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
def test_post_completions_privacy_mode(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def test_post_completions_returns_sse_done(monkeypatch):
|
||||
async def fake_call(*args, **kwargs):
|
||||
captured["kwargs"] = kwargs
|
||||
return {"content": "done", "think": ""}
|
||||
monkeypatch.setattr(main, "call_ollama", fake_call)
|
||||
monkeypatch.setattr(main, "build_completion_prompts", lambda *a, **k: ("sys", "user", ""))
|
||||
monkeypatch.setattr(main, "prepare_prompt_context", lambda *a, **k: ("p", "s"))
|
||||
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions", headers=HEADERS, json={
|
||||
"prefix": "hello", "suffix": "", "languageId": "markdown",
|
||||
"model_thinking": "low", "privacy_mode": True,
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data.get("content") == "done"
|
||||
# enable_thinking removed in OpenAI-compatible rewrite
|
||||
assert captured["kwargs"]["thinking"] == "low"
|
||||
|
||||
|
||||
def test_old_post_pro_stream_returns_404():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/pro/completions/stream", headers=HEADERS, json={
|
||||
"prefix": "hello",
|
||||
"suffix": "",
|
||||
"languageId": "markdown",
|
||||
"model_thinking": "high",
|
||||
"privacy_mode": True,
|
||||
})
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
def test_post_pro_completion_returns_sse_and_status(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
async def fake_stream_events(*args, **kwargs):
|
||||
captured["kwargs"] = kwargs
|
||||
yield "thinking", ""
|
||||
yield "content", "深度"
|
||||
yield "content", "回答"
|
||||
|
||||
monkeypatch.setattr(pro_completions, "stream_ollama_events", fake_stream_events)
|
||||
client = TestClient(main.app)
|
||||
with client.stream("POST", "/v1/pro/completions", headers=HEADERS, json={
|
||||
"prefix": "hello",
|
||||
"suffix": "",
|
||||
"languageId": "markdown",
|
||||
"instruction": "expand",
|
||||
"pro_thinking": "high",
|
||||
"privacy_mode": True,
|
||||
}) as resp:
|
||||
assert resp.status_code == 200
|
||||
body = "".join(resp.iter_text())
|
||||
return {"content": "系统done", "think": ""}
|
||||
|
||||
monkeypatch.setattr(job_handlers, "call_ollama", fake_call)
|
||||
with TestClient(main.app) as client:
|
||||
with client.stream("POST", "/v1/completions", headers=HEADERS, json={
|
||||
"prefix": "hello", "suffix": "", "languageId": "markdown",
|
||||
"model_thinking": "low", "privacy_mode": True,
|
||||
}) as resp:
|
||||
assert resp.status_code == 200
|
||||
body = "".join(resp.iter_text())
|
||||
assert "event: queued" in body
|
||||
assert "event: started" in body
|
||||
assert "event: thinking" in body
|
||||
assert "event: chunk" in body
|
||||
assert "event: result" in body
|
||||
assert "event: done" in body
|
||||
assert "深度" in body
|
||||
assert "回答" in body
|
||||
assert captured["kwargs"]["use_pro_model"] is True
|
||||
assert captured["kwargs"]["thinking"] == "high"
|
||||
|
||||
request_id = next(iter(pro_completions.PRO_STATES))
|
||||
status_resp = client.get(f"/v1/pro/completions/status/{request_id}", headers=HEADERS)
|
||||
assert status_resp.status_code == 200
|
||||
assert status_resp.json()["status"] == "done"
|
||||
assert main.ACTIVE_COMPLETIONS == {}
|
||||
|
||||
|
||||
def test_post_ocr_mocked(monkeypatch):
|
||||
async def fake_ocr(*args, **kwargs):
|
||||
return "OCR result text"
|
||||
monkeypatch.setattr(main, "call_vlm_ocr", fake_ocr)
|
||||
|
||||
client = TestClient(main.app)
|
||||
monkeypatch.setattr(job_handlers, "call_vlm_ocr", fake_ocr)
|
||||
img_b64 = base64.b64encode(b"pretend image data").decode()
|
||||
resp = client.post("/v1/ocr", headers=HEADERS, json={
|
||||
"image": img_b64, "filename": "test.jpg", "language": "auto",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
j = resp.json()
|
||||
assert j["text"] == "OCR result text"
|
||||
assert j["filename"] == "test.jpg"
|
||||
|
||||
|
||||
def test_post_ocr_invalid_base64_returns_500():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/ocr", headers=HEADERS, json={
|
||||
"image": "not-base64!!!", "filename": "test.jpg",
|
||||
})
|
||||
assert resp.status_code == 500
|
||||
with TestClient(main.app) as client:
|
||||
with client.stream("POST", "/v1/ocr", headers=HEADERS, json={
|
||||
"image": img_b64, "filename": "test.jpg", "language": "auto",
|
||||
}) as resp:
|
||||
assert resp.status_code == 200
|
||||
body = "".join(resp.iter_text())
|
||||
assert "OCR result text" in body
|
||||
|
||||
|
||||
def test_post_convert_txt_returns_markdown():
|
||||
client = TestClient(main.app)
|
||||
content = base64.b64encode(b"hello world").decode()
|
||||
resp = client.post("/v1/convert", headers=HEADERS, json={
|
||||
"file": content, "filename": "sample.txt",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
j = resp.json()
|
||||
assert j["markdown"] == "hello world"
|
||||
assert j["filename"] == "sample.txt"
|
||||
with TestClient(main.app) as client:
|
||||
with client.stream("POST", "/v1/convert", headers=HEADERS, json={
|
||||
"file": content, "filename": "sample.txt",
|
||||
}) as resp:
|
||||
assert resp.status_code == 200
|
||||
body = "".join(resp.iter_text())
|
||||
assert "hello world" in body
|
||||
|
||||
|
||||
def test_post_convert_unsupported_extension_returns_500():
|
||||
client = TestClient(main.app)
|
||||
content = base64.b64encode(b"data").decode()
|
||||
resp = client.post("/v1/convert", headers=HEADERS, json={
|
||||
"file": content, "filename": "sample.xlsx",
|
||||
})
|
||||
with TestClient(main.app) as client:
|
||||
resp = client.post("/v1/convert", headers=HEADERS, json={
|
||||
"file": content, "filename": "sample.xlsx",
|
||||
})
|
||||
assert resp.status_code == 500
|
||||
assert "仅支持" in resp.json()["error"]
|
||||
|
||||
|
||||
def test_post_convert_docx_with_mocked_markitdown(monkeypatch):
|
||||
class FakeResult:
|
||||
text_content = "markdown from docx"
|
||||
class FakeMD:
|
||||
def convert(self, path):
|
||||
return FakeResult()
|
||||
monkeypatch.setattr(main, "_get_markitdown", lambda: FakeMD())
|
||||
|
||||
client = TestClient(main.app)
|
||||
content = base64.b64encode(b"docx content").decode()
|
||||
resp = client.post("/v1/convert", headers=HEADERS, json={
|
||||
"file": content, "filename": "sample.docx",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
j = resp.json()
|
||||
assert j["markdown"] == "markdown from docx"
|
||||
|
||||
|
||||
def test_post_cancel_non_existent_returns_not_found():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions/cancel", headers=HEADERS, json={
|
||||
"request_id": "non-existent", "reason": "abort",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["cancelled"] is False
|
||||
assert data["status"] == "not_found"
|
||||
|
||||
|
||||
def test_post_cancel_wrong_api_key_returns_401():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions/cancel", json={
|
||||
"request_id": "id", "reason": "abort",
|
||||
})
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
def test_post_cancel_already_done(monkeypatch):
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
# Create a mock task that appears done
|
||||
mock_task = MagicMock()
|
||||
mock_task.done.return_value = True
|
||||
mock_task.cancel = MagicMock()
|
||||
main.ACTIVE_COMPLETIONS["done-id"] = mock_task
|
||||
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions/cancel", headers=HEADERS, json={
|
||||
"request_id": "done-id", "reason": "abort",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["cancelled"] is False
|
||||
assert data["status"] == "already_done"
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
|
||||
@@ -1,28 +1,30 @@
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
import asyncio
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
os.environ["JOB_BACKEND"] = "memory"
|
||||
|
||||
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
BACKEND_DIR = os.path.abspath(os.path.join(CURRENT_DIR, ".."))
|
||||
if BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, BACKEND_DIR)
|
||||
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(BACKEND_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(BACKEND_DIR))
|
||||
|
||||
if "tts_asr" not in sys.modules:
|
||||
fake_tts_asr = types.ModuleType("tts_asr")
|
||||
fake_tts_asr.register_tts_asr_routes = lambda app: None
|
||||
sys.modules["tts_asr"] = fake_tts_asr
|
||||
|
||||
import main # type: ignore
|
||||
import pro_completions # type: ignore
|
||||
import job_handlers # type: ignore
|
||||
import job_system # type: ignore
|
||||
import prompt # type: ignore
|
||||
|
||||
main = importlib.import_module("main")
|
||||
|
||||
HEADERS = {"X-API-Key": main.API_KEY}
|
||||
|
||||
|
||||
def setup_function():
|
||||
job_system.reset_job_manager()
|
||||
main._handlers_registered = False
|
||||
|
||||
|
||||
def _payload():
|
||||
return {
|
||||
"prefix": "Before",
|
||||
@@ -34,84 +36,52 @@ def _payload():
|
||||
}
|
||||
|
||||
|
||||
def setup_function():
|
||||
pro_completions.PRO_STATES.clear()
|
||||
|
||||
|
||||
def teardown_function():
|
||||
pro_completions.PRO_STATES.clear()
|
||||
|
||||
|
||||
def test_pro_queue_full_returns_429(monkeypatch):
|
||||
monkeypatch.setattr(pro_completions, "PRO_QUEUE_MAX_SIZE", 0)
|
||||
client = TestClient(main.app)
|
||||
response = client.post("/v1/pro/completions", headers=HEADERS, json=_payload())
|
||||
async def fake_queue_job(*args, **kwargs):
|
||||
raise job_system.QueueFullError("pro_completion queue is full")
|
||||
|
||||
monkeypatch.setattr(main, "_queue_job", fake_queue_job)
|
||||
with TestClient(main.app) as client:
|
||||
response = client.post("/v1/pro/completions", headers=HEADERS, json=_payload())
|
||||
assert response.status_code == 429
|
||||
assert response.json()["error"] == "PRO queue is full"
|
||||
|
||||
|
||||
def test_pro_status_missing_returns_404():
|
||||
client = TestClient(main.app)
|
||||
response = client.get("/v1/pro/completions/status/missing", headers=HEADERS)
|
||||
with TestClient(main.app) as client:
|
||||
response = client.get("/v1/pro/completions/status/missing", headers=HEADERS)
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
def test_pro_prompt_uses_pro_specific_instruction():
|
||||
system_prompt, user_prompt = pro_completions._build_pro_prompts(
|
||||
system_prompt, user_prompt = prompt.build_pro_completion_prompts(
|
||||
prefix="欢迎使用 LLM-IN-TEXT\n\n即时可用的 LLM 系统",
|
||||
suffix="",
|
||||
language_id="markdown",
|
||||
instruction="",
|
||||
pro_thinking="high",
|
||||
pro_thinking_level="high",
|
||||
)
|
||||
combined = f"{system_prompt}\n{user_prompt}".lower()
|
||||
assert "[pro] model for llm-in-text" in combined
|
||||
assert "pro_mode: true" in combined
|
||||
assert "pro_thinking_level: high" in combined
|
||||
assert "long paragraphs or section-level output are allowed" in combined
|
||||
assert "highest priority" in combined
|
||||
assert "never copy tags to output" in combined
|
||||
assert "write only the markdown that belongs at the cursor" not in combined
|
||||
assert "continue the markdown naturally" in combined
|
||||
|
||||
|
||||
def test_pro_cancel_waits_for_stream_cleanup(monkeypatch):
|
||||
started = threading.Event()
|
||||
cleaned = threading.Event()
|
||||
|
||||
def test_pro_stream_returns_standard_events(monkeypatch):
|
||||
async def fake_stream_events(*args, **kwargs):
|
||||
started.set()
|
||||
try:
|
||||
yield "thinking", ""
|
||||
while True:
|
||||
await asyncio.sleep(0.05)
|
||||
finally:
|
||||
cleaned.set()
|
||||
|
||||
monkeypatch.setattr(pro_completions, "stream_ollama_events", fake_stream_events)
|
||||
request_id = "pro-cancel-cleanup"
|
||||
headers = {**HEADERS, "X-Request-Id": request_id}
|
||||
response_box = {}
|
||||
yield "thinking", ""
|
||||
yield "content", "深度"
|
||||
yield "content", "回答"
|
||||
|
||||
monkeypatch.setattr(job_handlers, "stream_ollama_events", fake_stream_events)
|
||||
with TestClient(main.app) as client:
|
||||
def send_stream():
|
||||
with client.stream("POST", "/v1/pro/completions", headers=headers, json=_payload()) as response:
|
||||
response_box["status_code"] = response.status_code
|
||||
response_box["body"] = "".join(response.iter_text())
|
||||
with client.stream("POST", "/v1/pro/completions", headers=HEADERS, json=_payload()) as resp:
|
||||
assert resp.status_code == 200
|
||||
body = "".join(resp.iter_text())
|
||||
|
||||
stream_thread = threading.Thread(target=send_stream, daemon=True)
|
||||
stream_thread.start()
|
||||
|
||||
assert started.wait(timeout=2.0)
|
||||
cancel_response = client.post(
|
||||
"/v1/pro/completions/cancel",
|
||||
headers=HEADERS,
|
||||
json={"request_id": request_id, "reason": "test"},
|
||||
)
|
||||
|
||||
assert cancel_response.status_code == 200
|
||||
assert cancel_response.json() == {"cancelled": True, "status": "ok"}
|
||||
assert cleaned.wait(timeout=2.0)
|
||||
|
||||
stream_thread.join(timeout=5.0)
|
||||
assert not stream_thread.is_alive()
|
||||
assert "event: queued" in body
|
||||
assert "event: started" in body
|
||||
assert "event: progress" in body
|
||||
assert "event: result" in body
|
||||
assert "event: done" in body
|
||||
assert "深度" in body
|
||||
assert "回答" in body
|
||||
|
||||
+44
-48
@@ -40,7 +40,8 @@ try:
|
||||
except Exception as e: # pragma: no cover
|
||||
logger.debug("modelscope import failed (optional): %s", e)
|
||||
|
||||
router = APIRouter()
|
||||
meta_router = APIRouter()
|
||||
generation_router = APIRouter()
|
||||
|
||||
# Global model instances
|
||||
_tts_model: Optional["Qwen3TTSModel"] = None
|
||||
@@ -314,7 +315,7 @@ def _ensure_align_model():
|
||||
return _align_model
|
||||
|
||||
|
||||
@router.get("/status", response_model=ModelStatus)
|
||||
@meta_router.get("/status", response_model=ModelStatus)
|
||||
async def get_status():
|
||||
"""获取模型状态"""
|
||||
return ModelStatus(
|
||||
@@ -324,7 +325,7 @@ async def get_status():
|
||||
)
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
@meta_router.get("/config")
|
||||
async def get_config():
|
||||
"""获取配置信息"""
|
||||
return {
|
||||
@@ -340,7 +341,7 @@ async def get_config():
|
||||
}
|
||||
|
||||
|
||||
@router.post("/warmup")
|
||||
@meta_router.post("/warmup")
|
||||
async def warmup_models():
|
||||
"""手动触发模型预热"""
|
||||
await _warmup_tts()
|
||||
@@ -355,39 +356,34 @@ async def warmup_models():
|
||||
}
|
||||
|
||||
|
||||
@router.post("/tts", response_model=TTSResponse)
|
||||
async def tts_endpoint(req: TTSRequest):
|
||||
"""TTS 文字转语音端点"""
|
||||
async def generate_tts_response(
|
||||
text: str,
|
||||
instruct: str = "",
|
||||
speaker: str = "Vivian",
|
||||
output_format: str = "wav",
|
||||
) -> TTSResponse:
|
||||
del speaker # current model path does not expose multi-speaker routing
|
||||
del output_format # current implementation always returns wav
|
||||
try:
|
||||
model = _ensure_tts_model()
|
||||
except Exception as e: # noqa: ANN001
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
text = req.text
|
||||
instruct = req.instruct or ""
|
||||
|
||||
try:
|
||||
# VoiceDesign 模型使用 generate_voice_design 方法
|
||||
wavs, sr = model.generate_voice_design( # type: ignore
|
||||
text=text,
|
||||
language="Chinese",
|
||||
instruct=instruct,
|
||||
instruct=instruct or "",
|
||||
)
|
||||
except Exception as e: # noqa: ANN001
|
||||
logger.exception("TTS 推理失败")
|
||||
raise HTTPException(status_code=500, detail=f"TTS 推理失败: {e}")
|
||||
|
||||
# Get first audio data
|
||||
wav_data = wavs[0] if isinstance(wavs, (list, tuple)) else wavs
|
||||
|
||||
# Convert to numpy array
|
||||
if hasattr(wav_data, 'numpy'): # type: ignore
|
||||
wav_data = wav_data.cpu().numpy() # type: ignore
|
||||
wav_data = np.asarray(wav_data, dtype=np.float32)
|
||||
|
||||
logger.debug("wav_data shape: %s, dtype: %s, sr: %s", wav_data.shape, wav_data.dtype, sr)
|
||||
|
||||
# Encode WAV to memory
|
||||
tmp_path = None
|
||||
try:
|
||||
import soundfile as sf # type: ignore
|
||||
@@ -404,22 +400,15 @@ async def tts_endpoint(req: TTSRequest):
|
||||
if tmp_path and os.path.exists(tmp_path): # noqa: SIM201
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except Exception as e: # noqa: ANN001
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
duration_ms = int(len(wav_data) / sr * 1000) if sr > 0 else 0
|
||||
|
||||
audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
|
||||
return TTSResponse(
|
||||
audio_base64=audio_base64,
|
||||
format="wav",
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
return TTSResponse(audio_base64=audio_base64, format="wav", duration_ms=duration_ms)
|
||||
|
||||
|
||||
@router.post("/asr", response_model=ASRResponse)
|
||||
async def asr_endpoint(req: ASRRequest):
|
||||
"""语音识别端点(非流式)"""
|
||||
async def generate_asr_response(audio_bytes: bytes, language: Optional[str] = "zh-CN") -> ASRResponse:
|
||||
if Qwen3ASRModel is None:
|
||||
raise HTTPException(status_code=501, detail="mlx_audio 未安装,ASR 功能不可用")
|
||||
|
||||
@@ -429,10 +418,6 @@ async def asr_endpoint(req: ASRRequest):
|
||||
raise HTTPException(status_code=500, detail=f"ASR 模型加载失败: {e}")
|
||||
|
||||
try:
|
||||
# Decode base64 audio to WAV bytes
|
||||
audio_bytes = base64.b64decode(req.audio_base64)
|
||||
|
||||
# Load WAV file and convert to 16kHz mono numpy array
|
||||
wav_buffer = io.BytesIO(audio_bytes)
|
||||
with wave.open(wav_buffer, 'rb') as wf: # noqa: SIM115
|
||||
n_channels = wf.getnchannels()
|
||||
@@ -443,11 +428,9 @@ async def asr_endpoint(req: ASRRequest):
|
||||
raw_data = wf.readframes(n_frames)
|
||||
audio_array = np.frombuffer(raw_data, dtype=np.int16 if sampwidth == 2 else np.float32)
|
||||
|
||||
# Convert to mono
|
||||
if n_channels > 1:
|
||||
audio_array = np.mean(audio_array.reshape(-1, n_channels), axis=1)
|
||||
|
||||
# Resample to 16kHz if needed
|
||||
if framerate != 16000:
|
||||
try:
|
||||
import scipy.signal as signal # type: ignore
|
||||
@@ -457,34 +440,47 @@ async def asr_endpoint(req: ASRRequest):
|
||||
except Exception as e2: # noqa: ANN001
|
||||
logger.warning("重采样失败,使用原始音频: %s", e2)
|
||||
|
||||
# Convert to float32 normalized
|
||||
if audio_array.dtype == np.int16:
|
||||
audio_array = audio_array.astype(np.float32) / 32768.0
|
||||
|
||||
# Run ASR inference (non-streaming)
|
||||
result = model.generate( # type: ignore
|
||||
audio_array,
|
||||
language=req.language if req.language else None,
|
||||
language=language if language else None,
|
||||
)
|
||||
|
||||
# Extract text and detected language from result (STTOutput)
|
||||
recognized_text = getattr(result, 'text', str(result)) if hasattr(result, 'text') else str(result)
|
||||
detected_lang = getattr(result, 'language', req.language or "zh-CN")
|
||||
|
||||
# If language is a list (from segments), take the first one
|
||||
detected_lang = getattr(result, 'language', language or "zh-CN")
|
||||
if isinstance(detected_lang, list) and len(detected_lang) > 0:
|
||||
detected_lang = detected_lang[0]
|
||||
|
||||
return ASRResponse(
|
||||
text=recognized_text,
|
||||
language=str(detected_lang),
|
||||
)
|
||||
|
||||
return ASRResponse(text=recognized_text, language=str(detected_lang))
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e: # noqa: ANN001
|
||||
logger.exception("ASR 推理失败")
|
||||
raise HTTPException(status_code=500, detail=f"ASR 推理失败: {e}")
|
||||
|
||||
|
||||
def register_tts_asr_routes(app):
|
||||
@generation_router.post("/tts", response_model=TTSResponse)
|
||||
async def tts_endpoint(req: TTSRequest):
|
||||
"""TTS 文字转语音端点"""
|
||||
return await generate_tts_response(
|
||||
text=req.text,
|
||||
instruct=req.instruct or "",
|
||||
speaker=req.speaker,
|
||||
output_format=req.format,
|
||||
)
|
||||
|
||||
|
||||
@generation_router.post("/asr", response_model=ASRResponse)
|
||||
async def asr_endpoint(req: ASRRequest):
|
||||
"""语音识别端点(非流式)"""
|
||||
audio_bytes = base64.b64decode(req.audio_base64)
|
||||
return await generate_asr_response(audio_bytes, req.language if req.language else None)
|
||||
|
||||
|
||||
def register_tts_asr_routes(app, include_generation_routes: bool = True):
|
||||
"""注册 TTS/ASR 路由到 FastAPI 应用"""
|
||||
app.include_router(router, prefix="/v1/tts-asr")
|
||||
app.include_router(meta_router, prefix="/v1/tts-asr")
|
||||
if include_generation_routes:
|
||||
app.include_router(generation_router, prefix="/v1/tts-asr")
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from job_handlers import (
|
||||
asr_handler,
|
||||
completion_handler,
|
||||
compress_handler,
|
||||
convert_handler,
|
||||
ocr_handler,
|
||||
pro_completion_handler,
|
||||
tts_handler,
|
||||
)
|
||||
from job_system import RedisJobManager, RedisWorker
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(name)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger("worker")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
manager = RedisJobManager()
|
||||
manager.register_handler("completion", completion_handler)
|
||||
manager.register_handler("pro_completion", pro_completion_handler)
|
||||
manager.register_handler("compress", compress_handler)
|
||||
manager.register_handler("ocr", ocr_handler)
|
||||
manager.register_handler("convert", convert_handler)
|
||||
manager.register_handler("tts", tts_handler)
|
||||
manager.register_handler("asr", asr_handler)
|
||||
worker = RedisWorker(manager)
|
||||
try:
|
||||
await worker.run_forever()
|
||||
finally:
|
||||
await manager.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user