Migrate backend jobs to Redis Streams

This commit is contained in:
“ydy0615”
2026-06-06 15:44:00 +08:00
parent 2c7a02f587
commit 81f711ef0b
69 changed files with 2709 additions and 6291 deletions
+35 -5
View File
@@ -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
+256
View File
@@ -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)
+704
View File
@@ -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
View File
@@ -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
View File
@@ -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)
+3
View File
@@ -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 -1
View File
@@ -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
View File
@@ -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
+79
View File
@@ -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"}
+19 -45
View File
@@ -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 == {}
+54 -239
View File
@@ -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 "![alt](image.png)" not in main._sanitize_converted_markdown(
"text with image ![alt](image.png) 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 "![alt](image.png)" not in main._sanitize_converted_markdown("text ![alt](image.png)")
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()
+41 -71
View File
@@ -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
View File
@@ -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")
+39
View File
@@ -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())