From 4b08504ff789c2af4598ffd73af80c555806ab0c Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 26 Jan 2024 22:39:56 -0300 Subject: [PATCH] Add files router to API v1 --- src/backend/langflow/api/router.py | 2 + src/backend/langflow/api/v1/__init__.py | 2 + src/backend/langflow/api/v1/files.py | 55 +++++++++++-------------- 3 files changed, 29 insertions(+), 30 deletions(-) diff --git a/src/backend/langflow/api/router.py b/src/backend/langflow/api/router.py index 24d64b401..060de2ea7 100644 --- a/src/backend/langflow/api/router.py +++ b/src/backend/langflow/api/router.py @@ -6,6 +6,7 @@ from langflow.api.v1 import ( chat_router, credentials_router, endpoints_router, + files_router, flows_router, login_router, store_router, @@ -25,3 +26,4 @@ router.include_router(users_router) router.include_router(api_key_router) router.include_router(login_router) router.include_router(credentials_router) +router.include_router(files_router) diff --git a/src/backend/langflow/api/v1/__init__.py b/src/backend/langflow/api/v1/__init__.py index 6368a0cc8..f42d3164a 100644 --- a/src/backend/langflow/api/v1/__init__.py +++ b/src/backend/langflow/api/v1/__init__.py @@ -2,6 +2,7 @@ from langflow.api.v1.api_key import router as api_key_router from langflow.api.v1.chat import router as chat_router from langflow.api.v1.credential import router as credentials_router from langflow.api.v1.endpoints import router as endpoints_router +from langflow.api.v1.files import router as files_router from langflow.api.v1.flows import router as flows_router from langflow.api.v1.login import router as login_router from langflow.api.v1.store import router as store_router @@ -18,4 +19,5 @@ __all__ = [ "api_key_router", "login_router", "credentials_router", + "files_router", ] diff --git a/src/backend/langflow/api/v1/files.py b/src/backend/langflow/api/v1/files.py index fca42bf61..d5c43db21 100644 --- a/src/backend/langflow/api/v1/files.py +++ b/src/backend/langflow/api/v1/files.py @@ -1,37 +1,22 @@ import hashlib from http import HTTPStatus -from pathlib import Path +from io import BytesIO from fastapi import APIRouter, Depends, HTTPException, UploadFile +from fastapi.responses import StreamingResponse from langflow.services.deps import get_storage_service from langflow.services.storage.service import StorageService +from langflow.services.storage.utils import build_content_type_from_extension router = APIRouter(tags=["Files"], prefix="/files") -@router.post("/upload/{folder}", status_code=HTTPStatus.CREATED) +@router.post("/upload/{flow_id}", status_code=HTTPStatus.CREATED) async def upload_file(flow_id: str, file: UploadFile, storage_service: StorageService = Depends(get_storage_service)): try: file_content = await file.read() - file_name = file.filename - filename = file.filename - if isinstance(filename, str) or isinstance(filename, Path): - file_extension = Path(filename).suffix - else: - file_extension = "" - file_object = file.file - sha256_hash = hashlib.sha256() - # Reset the file cursor to the beginning of the file - file_object.seek(0) - # Iterate over the uploaded file in small chunks to conserve memory - while chunk := file_object.read(8192): # Read 8KB at a time (adjust as needed) - sha256_hash.update(chunk) - - # Use the hex digest of the hash as the file name - hex_dig = sha256_hash.hexdigest() - file_name = f"{hex_dig}{file_extension}" - + file_name = file.filename or hashlib.sha256(file_content).hexdigest() folder = flow_id storage_service.save_file(folder=folder, file_name=file_name, data=file_content) return {"message": "File uploaded successfully", "file_path": f"{folder}/{file.filename}"} @@ -39,28 +24,38 @@ async def upload_file(flow_id: str, file: UploadFile, storage_service: StorageSe raise HTTPException(status_code=500, detail=str(e)) -@router.get("/download/{folder}/{file_name}") -async def download_file(folder: str, file_name: str, storage_service: StorageService = Depends(get_storage_service)): +@router.get("/download/{flow_id}/{file_name}") +async def download_file(flow_id: str, file_name: str, storage_service: StorageService = Depends(get_storage_service)): try: - file_content = storage_service.get_file(folder=folder, file_name=file_name) - return {"file_content": file_content} + extension = file_name.split(".")[-1] + + if not extension: + raise HTTPException(status_code=500, detail=f"Extension not found for file {file_name}") + + content_type = build_content_type_from_extension(extension) + + if not content_type: + raise HTTPException(status_code=500, detail=f"Content type not found for extension {extension}") + + file_content = storage_service.get_file(folder=flow_id, file_name=file_name) + return StreamingResponse(BytesIO(file_content), media_type=content_type) except Exception as e: raise HTTPException(status_code=500, detail=str(e)) -@router.get("/list/{folder}") -async def list_files(folder: str, storage_service: StorageService = Depends(get_storage_service)): +@router.get("/list/{flow_id}") +async def list_files(flow_id: str, storage_service: StorageService = Depends(get_storage_service)): try: - files = storage_service.list_files(folder=folder) + files = storage_service.list_files(folder=flow_id) return {"files": files} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) -@router.delete("/delete/{folder}/{file_name}") -async def delete_file(folder: str, file_name: str, storage_service: StorageService = Depends(get_storage_service)): +@router.delete("/delete/{flow_id}/{file_name}") +async def delete_file(flow_id: str, file_name: str, storage_service: StorageService = Depends(get_storage_service)): try: - storage_service.delete_file(folder=folder, file_name=file_name) + storage_service.delete_file(folder=flow_id, file_name=file_name) return {"message": f"File {file_name} deleted successfully"} except Exception as e: raise HTTPException(status_code=500, detail=str(e))