llm-in-text/backend/main.py

import asyncio
import base64
import json
import logging
import os
import uuid
from typing import Optional

from fastapi import FastAPI, File, Form, HTTPException, Request, Response, Security, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.security import APIKeyHeader
from pydantic import BaseModel

from docs_store import get_document_store
from geoip import get_ip_location_text
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

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s - %(message)s",
)
logger = logging.getLogger("api")

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*", "X-API-Key", "X-Client-IP", "X-Request-Id"],
)

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")
_handlers_registered = False


class CompletionRequest(BaseModel):
    prefix: str
    suffix: str
    languageId: str = "markdown"
    model_thinking: str = "low"
    privacy_mode: bool = False
    user_preferences: Optional[UserPreferences] = None
    model: Optional[str] = None
    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"


class OCRRequest(BaseModel):
    image: str
    filename: str = "image.jpg"
    language: str = "auto"


class ConvertRequest(BaseModel):
    file: str
    filename: str = "document.pdf"


class CompressRequest(BaseModel):
    content: str
    docType: str = "txt"


class TTSJobRequest(BaseModel):
    text: str
    instruct: str = ""
    speaker: str = "Vivian"
    format: str = "wav"


class ASRJobRequest(BaseModel):
    audio_base64: str
    language: Optional[str] = "zh-CN"


class CreateFolderRequest(BaseModel):
    name: str
    parentId: Optional[str] = None


class CreateTextFileRequest(BaseModel):
    name: str
    parentId: Optional[str] = None
    content: str = ""


class UpdateNodeRequest(BaseModel):
    name: Optional[str] = None
    parentId: Optional[str] = None
    content: Optional[str] = None


def _preview(text: str, limit: int = 80) -> str:
    value = (text or "").replace("\n", "\\n")
    if len(value) <= limit:
        return value
    return value[:limit] + "..."


def get_client_ip(request: Request) -> str:
    if request.client:
        return request.headers.get("X-Client-IP") or request.client.host
    return request.headers.get("X-Client-IP") or "unknown"


def _clamp_temperature(value: float, default: float = 0.7) -> float:
    try:
        numeric = float(value)
    except (TypeError, ValueError):
        return default
    return max(0.0, min(numeric, 1.2))


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


def _serialize_preferences(preferences: UserPreferences | None) -> dict | None:
    if preferences is None:
        return None
    if hasattr(preferences, "dict"):
        return preferences.dict()
    return dict(preferences)


def _request_id(request: Request) -> str:
    return request.headers.get("X-Request-Id") or str(uuid.uuid4())


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


def _sse(event: str, data: dict) -> str:
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"


async def _stream_job(job_id: str):
    _register_handlers()
    manager = get_job_manager()

    async def event_stream():
        try:
            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(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


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 {}


async def _docs_store_call(method_name: str, *args, **kwargs):
    store = get_document_store()
    method = getattr(store, method_name)
    return await asyncio.to_thread(method, *args, **kwargs)


@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)):
    del api_key
    return await _cancel_job(req.request_id or "", req.reason)


@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)


@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(req: OCRRequest, api_key: str = Security(get_api_key)):
    del api_key
    request_id = str(uuid.uuid4())
    try:
        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(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:
        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"}


@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"),
    }


@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)


@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)


@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()}


@app.get("/v1/docs/nodes")
async def list_docs_nodes(api_key: str = Security(get_api_key)):
    del api_key
    try:
        return {"nodes": await _docs_store_call("list_nodes")}
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc


@app.post("/v1/docs/folders")
async def create_docs_folder(req: CreateFolderRequest, api_key: str = Security(get_api_key)):
    del api_key
    if not (req.name or "").strip():
        raise HTTPException(status_code=400, detail="文件夹名称不能为空")
    try:
        node = await _docs_store_call("create_folder", req.name.strip(), req.parentId)
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc
    return {"node": node}


@app.post("/v1/docs/files/text")
async def create_docs_text_file(req: CreateTextFileRequest, api_key: str = Security(get_api_key)):
    del api_key
    if not (req.name or "").strip():
        raise HTTPException(status_code=400, detail="文件名称不能为空")
    try:
        node = await _docs_store_call("create_text_file", req.name.strip(), req.parentId, req.content or "")
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc
    return {"node": node}


@app.post("/v1/docs/files/upload")
async def upload_docs_file(
    file: UploadFile = File(...),
    parent_id: Optional[str] = Form(default=None),
    api_key: str = Security(get_api_key),
):
    del api_key
    filename = (file.filename or "").strip()
    if not filename:
        raise HTTPException(status_code=400, detail="文件名称不能为空")
    raw_bytes = await file.read()
    try:
        node = await _docs_store_call("upload_file", filename, parent_id, raw_bytes, file.content_type or "")
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc
    return {"node": node}


@app.patch("/v1/docs/nodes/{node_id}")
async def update_docs_node(node_id: str, req: UpdateNodeRequest, api_key: str = Security(get_api_key)):
    del api_key
    fields_set = req.model_fields_set if hasattr(req, "model_fields_set") else getattr(req, "__fields_set__", set())
    if not fields_set:
        raise HTTPException(status_code=400, detail="缺少更新内容")
    update_kwargs = {}
    if "name" in fields_set:
        next_name = req.name.strip() if isinstance(req.name, str) else ""
        if not next_name:
            raise HTTPException(status_code=400, detail="名称不能为空")
        update_kwargs["name"] = next_name
    if "parentId" in fields_set:
        update_kwargs["parent_id"] = req.parentId
    if "content" in fields_set:
        update_kwargs["content"] = req.content or ""
    try:
        node = await _docs_store_call("update_node", node_id, **update_kwargs)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="节点不存在") from exc
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc
    return {"node": node}


@app.put("/v1/docs/files/{node_id}/blob")
async def replace_docs_blob(
    node_id: str,
    file: UploadFile = File(...),
    api_key: str = Security(get_api_key),
):
    del api_key
    filename = (file.filename or "").strip()
    if not filename:
        raise HTTPException(status_code=400, detail="文件名称不能为空")
    raw_bytes = await file.read()
    try:
        node = await _docs_store_call("replace_blob", node_id, filename, raw_bytes, file.content_type or "")
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="节点不存在") from exc
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc
    return {"node": node}


@app.delete("/v1/docs/nodes/{node_id}")
async def delete_docs_node(node_id: str, api_key: str = Security(get_api_key)):
    del api_key
    try:
        await _docs_store_call("delete_node", node_id)
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc
    return {"ok": True}


@app.get("/v1/docs/files/{node_id}/blob")
async def download_docs_blob(node_id: str, api_key: str = Security(get_api_key)):
    del api_key
    try:
        payload = await _docs_store_call("get_blob", node_id)
    except FileNotFoundError as exc:
        raise HTTPException(status_code=404, detail="文件不存在") from exc
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc
    headers = {
        "Content-Disposition": f'inline; filename="{payload.filename}"',
    }
    return Response(content=payload.content, media_type=payload.mime_type, headers=headers)


def _register_tts_asr_routes():
    try:
        from tts_asr import register_tts_asr_routes
    except ModuleNotFoundError as exc:
        logger.warning("Skipping TTS/ASR route registration because a dependency is missing: %s", exc)
        return
    except Exception as exc:
        logger.warning("Skipping TTS/ASR route registration because import failed: %s", exc)
        return

    try:
        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)