import os
from fastapi import FastAPI
from transformers import pipeline

# 必须在所有导入前设置环境变量
os.environ["HF_HOME"] = "/app/.cache/huggingface"
os.environ["TRANSFORMERS_CACHE"] = "/app/.cache/huggingface"

app = FastAPI()

# 模型加载（带错误处理）
try:
    analyzer = pipeline(
        "text2text-generation",
        model="Salesforce/codet5-small",
        tokenizer="Salesforce/codet5-small"
    )
except Exception as e:
    raise RuntimeError(f"模型加载失败: {str(e)}")

@app.post("/analyze")
async def analyze(code: str):
    result = analyzer(
        f"Analyze code vulnerabilities:\n{code}",
        max_length=512,
        num_beams=5
    )[0]['generated_text']
    return {"result": result}


# 添加GET支持，测试模型
@app.get("/analyze")
async def analyze_get(code: str):  # 新增GET方法
    return await analyze(code)