feat: upload endpoint

This commit is contained in:
Timothy J. Baek 2024-06-05 13:25:07 -07:00
parent 17a44df0f3
commit d08a1573da

48
main.py
View File

@ -1,4 +1,4 @@
from fastapi import FastAPI, Request, Depends, status, HTTPException from fastapi import FastAPI, Request, Depends, status, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from fastapi.concurrency import run_in_threadpool from fastapi.concurrency import run_in_threadpool
@ -17,6 +17,7 @@ from concurrent.futures import ThreadPoolExecutor
from schemas import FilterForm, OpenAIChatCompletionForm from schemas import FilterForm, OpenAIChatCompletionForm
from urllib.parse import urlparse from urllib.parse import urlparse
import shutil
import aiohttp import aiohttp
import os import os
import importlib.util import importlib.util
@ -372,6 +373,51 @@ async def add_pipeline(
) )
@app.post("/v1/pipelines/upload")
@app.post("/pipelines/upload")
async def upload_pipeline(
file: UploadFile = File(...), user: str = Depends(get_current_user)
):
if user != API_KEY:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid API key",
)
file_ext = os.path.splitext(file.filename)[1]
if file_ext != ".py":
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Only Python files are allowed.",
)
try:
# Ensure the destination folder exists
os.makedirs(PIPELINES_DIR, exist_ok=True)
# Define the file path
file_path = os.path.join(PIPELINES_DIR, file.filename)
# Save the uploaded file to the specified directory
with open(file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
# Perform any necessary reload or processing
await reload()
return {
"status": True,
"detail": f"Pipeline uploaded successfully to {file_path}",
}
except HTTPException as e:
raise e
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=str(e),
)
class DeletePipelineForm(BaseModel): class DeletePipelineForm(BaseModel):
id: str id: str