from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, T5ForConditionalGeneration
import torch

app = FastAPI()

# 全局加载模型
tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-small")
model = T5ForConditionalGeneration.from_pretrained("Salesforce/codet5-small")

class CodeRequest(BaseModel):
    code: str
    max_length: int = 512

@app.post("/v1/analyze")
async def analyze_code(request: CodeRequest):
    try:
        # 构造提示词
        prompt = f"Analyze security vulnerabilities in this code:\n{request.code}"
        
        # 生成分析结果
        inputs = tokenizer(prompt, return_tensors="pt", 
                          max_length=512, truncation=True)
        outputs = model.generate(
            inputs.input_ids,
            max_length=request.max_length,
            num_beams=5,
            early_stopping=True
        )
        
        # 解码结果
        analysis = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return {
            "status": "success",
            "analysis": analysis,
            "model": "Salesforce/codet5-small"
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))