Migrate backend jobs to Redis Streams
This commit is contained in:
“ydy0615”
+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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user