Add image download endpoint

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-29 14:22:18 -03:00
commit d692ccd3ab

View file

@ -4,6 +4,7 @@ from io import BytesIO
from fastapi import APIRouter, Depends, HTTPException, UploadFile from fastapi import APIRouter, Depends, HTTPException, UploadFile
from fastapi.responses import StreamingResponse from fastapi.responses import StreamingResponse
from langflow.api.v1.schemas import UploadFileResponse from langflow.api.v1.schemas import UploadFileResponse
from langflow.services.auth.utils import get_current_active_user from langflow.services.auth.utils import get_current_active_user
from langflow.services.database.models.flow import Flow from langflow.services.database.models.flow import Flow
@ -73,6 +74,29 @@ async def download_file(
raise HTTPException(status_code=500, detail=str(e)) raise HTTPException(status_code=500, detail=str(e))
@router("/images/{flow_id}/{file_name}")
async def download_image(
file_name: str, flow_id: str = Depends(get_flow_id), storage_service: StorageService = Depends(get_storage_service)
):
try:
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}")
elif not content_type.startswith("image"):
raise HTTPException(status_code=500, detail=f"Content type {content_type} is not an image")
file_content = storage_service.get_file(flow_id=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/{flow_id}") @router.get("/list/{flow_id}")
async def list_files( async def list_files(
flow_id: str = Depends(get_flow_id), storage_service: StorageService = Depends(get_storage_service) flow_id: str = Depends(get_flow_id), storage_service: StorageService = Depends(get_storage_service)