from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
import os

# 正确加载模型（从缓存或下载）
classifier = pipeline(
    "text-classification",
    model="mrm8488/codebert-base-finetuned-detect-insecure-code"
)

app = FastAPI()

class CodeRequest(BaseModel):
    code: str  # 输入参数定义

@app.post("/detect")
async def detect_insecure_code(request: CodeRequest):
    try:
        # 直接传递代码字符串到分类器
        result = classifier(request.code)
        return {"status": "success", "result": result[0]}
    except Exception as e:
        return {"status": "error", "message": str(e)}