import os
from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.middleware.wsgi import WSGIMiddleware


from transformers import pipeline

app = FastAPI()  # 创建 FastAPI 应用

# 定义请求模型
class TextRequest(BaseModel):
    text: str

# 定义两个 API 路由处理函数
@app.post("/api/aaa")
async def api_aaa_post(request: TextRequest):
    result = request.text + 'aaa'
    return {"result": result}

# 定义两个 API 路由处理函数
@app.post("/aaa")
async def aaa(request: TextRequest):
    result = request.text + 'aaa'
    return {"result": result}


# 定义两个 API 路由处理函数
@app.get("/aaa")
async def api_aaa_get(request: TextRequest):
    result = request.text + 'aaa'
    return {"result": result}

@app.post("/api/bbb")
async def api_bbb(request: TextRequest):
    result = request.text + 'bbb'
    return {"result": result}


pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")

@app.get("/infer_t5")
def t5_get(input):
    output = pipe_flan(input)
    return {"output": output[0]["generated_text"]}


@app.post("/infer_t5")
def t5_post(input):
    output = pipe_flan(input)
    return {"output": output[0]["generated_text"]}

@app.get("/")
async def root():
    return {"message": "Welcome to the API. Use /api/aaa or /api/bbb for processing."}