coze_znjd/main.py

# -*- coding:utf-8 -*-
"""
@author: isaacqyang
@time: 2022/8/29
@desc:
"""

import argparse
import time
import traceback
import uuid

import uvicorn
from fastapi import FastAPI, Request, Path, Body
from starlette.responses import RedirectResponse

from commom import f_doc_export, get_logger, request_id_context
from entitys.response import BaseResponse

logger = get_logger()
app = FastAPI()


def doc_export(
        token: str = Path(title='doc token', description="文档的token"),
        request_id: str = Path(title='request id', description="请求ID作为文档名后缀"),
        data: object|None = Body(None, description="额外参数")

):
    try:
        # request_id = str(uuid.uuid4()) #request_id_context.get()
        # logger.info(f"Request ID: {request_id}, Token: {token}")
        word_downloda_url = f_doc_export(token, request_id, data)
        if word_downloda_url:
            return BaseResponse.ofSuccess(word_downloda_url)
        else:
            return BaseResponse.ofFailure("生成word失败")

    except Exception as msg:
        logger.error(traceback.format_exc())
        return BaseResponse.ofFailure(str(msg))


@app.middleware("http")
async def add_request_id_header(request: Request, call_next):
    request_id = request.headers.get("X-REQUEST-ID")
    if request_id is None or len(request_id) == 0:
        request_id = str(uuid.uuid4())
    request_id_context.set(request_id)
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-REQUEST-ID"] = request_id_context.get()
    response.headers["PROCESS-TIME"] = f"{process_time:.2f}"

    return response


async def document():
    return RedirectResponse(url="/docs")


def api_start(host, port):
    app.get("/", response_model=BaseResponse)(document)
    app.post("/znjd/doc/export/{token}/{request_id}", response_model=BaseResponse)(doc_export)
    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=18070)
    args = parser.parse_args()
    api_start(args.host, args.port)