Tiktok-Talent-Info/main.py

# from fastapi import FastAPI, Form, UploadFile
# from fastapi.responses import JSONResponse
# import asyncio

# app = FastAPI()

# @app.post("/api/text")
# async def text_query_endpoint(question: str = Form(...)):
#     """
#     API endpoint to process text input with the user's query.
#     """
#     from endpoints.text import text_query 
#     return await text_query(question=question)


# @app.post("/api/image")
# async def image_query_endpoint(file: UploadFile, question: str = Form(...)):
#     """
#     API endpoint to process an image with the user's query.
#     """
#     from endpoints.image import image_query 
#     return await image_query(file=file, question=question)


# @app.post("/api/video")
# async def video_query_endpoint(file: UploadFile, question: str = Form(...)):
#     """
#     API endpoint to process a video file with the user's query.
#     """
#     from endpoints.video import video_query  
#     return await video_query(file=file, question=question)


# if __name__ == "__main__":
#     import uvicorn
#     uvicorn.run("main:app", host="0.0.0.0", port=8002, reload=True, loop="uvloop")


from fastapi import FastAPI, Form, UploadFile
from fastapi.responses import JSONResponse
import shutil
import uuid
from tasks import text_query_task, image_query_task, video_query_task

app = FastAPI()

# @app.post("/api/text")
# async def text_query_endpoint(question: str = Form(...)):
#     task = text_query_task.apply_async(args=[question])
#     return JSONResponse({"task_id": task.id})

@app.post("/api/text")
async def text_query_endpoint(question: str = Form(...)):
    print(f"Received request: {question}") 
    task = text_query_task.apply_async(args=[question])
    print(f"Task sent: {task.id}")
    return JSONResponse({"task_id": task.id})

@app.post("/api/image")
async def image_query_endpoint(file: UploadFile, question: str = Form(...)):
    file_path = f"/tmp/{uuid.uuid4()}_{file.filename}"
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    
    task = image_query_task.apply_async(args=[file_path, question])
    return JSONResponse({"task_id": task.id})

@app.post("/api/video")
async def video_query_endpoint(file: UploadFile, question: str = Form(...)):
    file_path = f"/tmp/{uuid.uuid4()}_{file.filename}"
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    
    task = video_query_task.apply_async(args=[file_path, question])
    return JSONResponse({"task_id": task.id})

@app.get("/api/task/{task_id}")
async def get_task_result(task_id: str):
    from celery.result import AsyncResult
    result = AsyncResult(task_id)
    if result.ready():
        return JSONResponse({"status": "completed", "result": result.result})
    return JSONResponse({"status": "pending"})