BrandCultivation/api.py

from fastapi import FastAPI, HTTPException, Request, status
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from database.dao.mysql_dao import MySqlDao
from models import Recommend
import os
from pydantic import BaseModel
import uvicorn

app = FastAPI()
dao = MySqlDao()

# 添加全局异常处理器
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    return JSONResponse(
        status_code=status.HTTP_400_BAD_REQUEST,
        content={
            "code": 400,
            "msg": "请求参数错误",
            "data": {
                "detail": exc.errors(),
                "body": exc.body
            }
        },
    )

# 定义请求体
class RecommendRequest(BaseModel):
    city_uuid: str              # 城市id
    product_code: str           # 卷烟编码
    recall_cust_count: int      # 推荐的商户数量
    delivery_count: int              # 投放的品规数量
    
@app.post("/brandcultivation/api/v1/recommend")
def recommend(request: RecommendRequest):
    gbdtlr_model_path = os.path.join("./models/rank/weights", request.city_uuid, "gbdtlr_model.pkl")
    if not os.path.exists(gbdtlr_model_path):
        return {"code": 200, "msg": "model not defined", "data": {"recommendationInfo": "该城市的模型未训练，请先进行训练"}}
    
    # 初始化模型
    recommend_model = Recommend(request.city_uuid)
    
    # 判断该品规是否是新品规
    products_in_oreder = dao.get_product_from_order(request.city_uuid)["product_code"].unique().tolist()
    if request.product_code in products_in_oreder:
        recommend_list = recommend_model.get_recommend_list_by_gbdtlr(request.product_code, recall_count=request.recall_cust_count)
    else:
        recommend_list = recommend_model.get_recommend_list_by_item2vec(request.product_code, recall_count=request.recall_cust_count)
    recommend_data = recommend_model.get_recommend_and_delivery(recommend_list, delivery_count=request.delivery_count)
    request_data = []
    for index, data in enumerate(recommend_data):
        id = index + 1
        request_data.append(
            {
                "id": id,
                "cust_code": data["cust_code"],
                "recommend_score": data["recommend_score"],
                "delivery_count": data["delivery_count"]
            }
        )
        
    return {"code": 200, "msg": "success", "data": {"recommendationInfo": request_data}}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)