ensure we dont keep generating new images

autogpt_platform/backend/backend/server/v2/store/media.py

@@ -17,10 +17,11 @@
 async def check_media_exists(user_id: str, filename: str) -> str | None:
     """
     Check if a media file exists in storage for the given user.
+    Tries both images and videos directories.
 
     Args:
         user_id (str): ID of the user who uploaded the file
-        file_name (str): Name of the file to check
+        filename (str): Name of the file to check
 
     Returns:
         str | None: URL of the blob if it exists, None otherwise
@@ -29,21 +30,27 @@ async def check_media_exists(user_id: str, filename: str) -> str | None:
         settings = Settings()
         storage_client = storage.Client()
         bucket = storage_client.bucket(settings.config.media_gcs_bucket_name)
-
-        # Construct the full path including user_id
-        blob_path = f"users/{user_id}/images/{filename}"
-        blob = bucket.blob(blob_path)
-
-        if blob.exists():
-            return blob.public_url
-        return None
 
+        # Check images
+        image_path = f"users/{user_id}/images/{filename}"
+        image_blob = bucket.blob(image_path)
+        if image_blob.exists():
+            return image_blob.public_url
+
+        # Check videos
+        video_path = f"users/{user_id}/videos/{filename}"
+
+        video_blob = bucket.blob(video_path)
+        if video_blob.exists():
+            return video_blob.public_url
+
+        return None
     except Exception as e:
         logger.error(f"Error checking if media file exists: {str(e)}")
         return None
 
 
-async def upload_media(user_id: str, file: fastapi.UploadFile) -> str:
+async def upload_media(user_id: str, file: fastapi.UploadFile, use_file_name: bool = False) -> str:
 
     # Get file content for deeper validation
     try:
@@ -150,7 +157,10 @@ async def upload_media(user_id: str, file: fastapi.UploadFile) -> str:
         # Generate unique filename
         filename = file.filename or ""
         file_ext = os.path.splitext(filename)[1].lower()
-        unique_filename = f"{uuid.uuid4()}{file_ext}"
+        if use_file_name:
+            unique_filename = filename
+        else:
+            unique_filename = f"{uuid.uuid4()}{file_ext}"
 
         # Construct storage path
         media_type = "images" if content_type in ALLOWED_IMAGE_TYPES else "videos"

autogpt_platform/backend/backend/server/v2/store/routes.py

@@ -1,16 +1,16 @@
 import logging
 import typing
 
-import autogpt_libs.auth.depends
-import autogpt_libs.auth.middleware
 import fastapi
 import fastapi.responses
 
+import autogpt_libs.auth.depends
+import autogpt_libs.auth.middleware
+import backend.data.graph
 import backend.server.v2.store.db
+import backend.server.v2.store.image_gen
 import backend.server.v2.store.media
 import backend.server.v2.store.model
-import backend.data.graph
-import backend.server.v2.store.image_gen
 
 logger = logging.getLogger(__name__)
 
@@ -462,43 +462,38 @@ async def generate_image(
         user_id (str): ID of the authenticated user
 
     Returns:
-        str: URL of the generated image
-
-    Raises:
-        HTTPException: If there is an error generating the image
+        JSONResponse: JSON containing the URL of the generated image
     """
     try:
-
         agent = await backend.data.graph.get_graph(agent_id, user_id=user_id)
 
         if not agent:
             raise fastapi.HTTPException(
                 status_code=404, detail=f"Agent with ID {agent_id} not found"
             )
+        # Use .jpeg here since we are generating JPEG images
+        filename = f"agent_{agent_id}.jpeg"
 
-        existing_url = await backend.server.v2.store.media.check_media_exists(user_id, f"agent_{agent_id}.png")
+        existing_url = await backend.server.v2.store.media.check_media_exists(user_id, filename)
         if existing_url:
             logger.info(f"Using existing image for agent {agent_id}")
             return fastapi.responses.JSONResponse(content={"image_url": existing_url})
+        # Generate agent image as JPEG
+        image = await backend.server.v2.store.image_gen.generate_agent_image(agent=agent)
 
-        image = await backend.server.v2.store.image_gen.generate_agent_image(
-            agent=agent
-        )
-
-        # Create UploadFile and then set content_type
+        # Create UploadFile with the correct filename and content_type
         image_file = fastapi.UploadFile(
             file=image,
-            filename=f"agent_{agent_id}.jpeg"
+            filename=filename,
         )
 
         image_url = await backend.server.v2.store.media.upload_media(
-            user_id=user_id, file=image_file
+            user_id=user_id, file=image_file, use_file_name=True
         )
-
 
         return fastapi.responses.JSONResponse(content={"image_url": image_url})
     except Exception as e:
         logger.exception("Exception occurred whilst generating submission image")
         raise fastapi.HTTPException(
             status_code=500, detail=f"Failed to generate image: {str(e)}"
-        )
+        )