Merge pull request #13 from GoogleCloudPlatform/feature/frontend-workbench

Feature/workbench added

Author: MauroCominotti

backend/src/common/base_dto.py

@@ -25,6 +25,18 @@ class MimeTypeEnum(str, Enum):
     IMAGE_PNG = "image/png"
     VIDEO_MP4 = "video/mp4"
     AUDIO_WAV = "audio/wav"
+    AUDIO_MPEG = "audio/mpeg"
+    AUDIO_MP3 = "audio/mp3"
+    AUDIO_OGG = "audio/ogg"
+    AUDIO_WEBM = "audio/webm"
+
+
+class WildcardMimeTypeEnum(str, Enum):
+    """Wildcard MIME types for search filtering."""
+
+    IMAGE_WILDCARD = "image/*"
+    VIDEO_WILDCARD = "video/*"
+    AUDIO_WILDCARD = "audio/*"
 
 
 class GenerationModelEnum(str, Enum):
@@ -156,4 +168,4 @@ class BaseDto(BaseModel):
         extra="forbid",
         populate_by_name=True,
         from_attributes=True,
-    )
+    )

backend/src/galleries/dto/gallery_search_dto.py

@@ -12,19 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Optional, Union
 
 from pydantic import Field
 
-from src.common.base_dto import GenerationModelEnum, MimeTypeEnum
+from src.common.base_dto import GenerationModelEnum, MimeTypeEnum, WildcardMimeTypeEnum
 from src.common.dto.base_search_dto import BaseSearchDto
 from src.common.schema.media_item_model import JobStatusEnum
 from src.galleries.dto.gallery_response_dto import MediaItemResponse
 
 
 class GallerySearchDto(BaseSearchDto):
     user_email: Optional[str] = None
-    mime_type: Optional[MimeTypeEnum] = None
+    mime_type: Optional[Union[MimeTypeEnum, WildcardMimeTypeEnum]] = None
     model: Optional[GenerationModelEnum] = None
     status: Optional[JobStatusEnum] = None
     workspace_id: int = Field(

backend/src/images/repository/media_item_repository.py

@@ -45,7 +45,13 @@ async def query(
             query = query.where(self.model.user_email == search_dto.user_email)
 
         if search_dto.mime_type:
-            query = query.where(self.model.mime_type == search_dto.mime_type.value)
+            if search_dto.mime_type.value.endswith("/*"):
+                # Handle wildcard search (e.g., "image/*")
+                prefix = search_dto.mime_type.value[:-1]  # Remove the "*"
+                query = query.where(self.model.mime_type.like(f"{prefix}%"))
+            else:
+                # Handle exact match
+                query = query.where(self.model.mime_type == search_dto.mime_type.value)
 
         if search_dto.model:
             query = query.where(self.model.model == search_dto.model.value)

backend/src/source_assets/repository/source_asset_repository.py

@@ -63,13 +63,17 @@ async def query(
 
         # Apply filters
         if search_dto.mime_type:
-            if search_dto.mime_type.endswith("image/*"):
-                # Handle wildcard prefix search
-                # In SQL, we can use LIKE 'image/%'
-                # But the original code did "!=" "video/mp4" which is weird for "image/*"
-                # Let's assume we want to match "image/%"
+            if search_dto.mime_type == "image/*":
+                # Filter for all image types using SQL LIKE
                 query = query.where(self.model.mime_type.like("image/%"))
+            elif search_dto.mime_type == "video/*":
+                # Filter for all video types
+                query = query.where(self.model.mime_type.like("video/%"))
+            elif search_dto.mime_type == "audio/*":
+                # Filter for all audio types
+                query = query.where(self.model.mime_type.like("audio/%"))
             else:
+                # Exact match for specific mime types like "image/png"
                 query = query.where(self.model.mime_type == search_dto.mime_type)
 
         if target_user_id:
@@ -184,4 +188,4 @@ async def find_system_and_private_assets_by_types(
             .where(self.model.asset_type.in_([t.value for t in asset_types]))
         )
         assets = result.scalars().all()
-        return [self.schema.model_validate(asset) for asset in assets]
+        return [self.schema.model_validate(asset) for asset in assets]

backend/src/source_assets/source_asset_service.py

@@ -189,10 +189,13 @@ async def upload_asset(
             )
             return await self._create_asset_response(existing_asset)
 
-        # 2. Handle file processing based on type (image vs. video)
+        # 2. Handle file processing based on type (image vs. video vs. audio)
         is_video: bool = bool(
             file.content_type and "video" in file.content_type
         )
+        is_audio: bool = bool(
+            file.content_type and "audio" in file.content_type
+        )
         final_gcs_uri: Optional[str] = None
         thumbnail_gcs_uri: Optional[str] = None
         temp_dir = f"temp/source_assets/{uuid.uuid4()}"
@@ -230,6 +233,23 @@ async def upload_asset(
                         destination_blob_name=f"source_assets/{user.id}/{file_hash}/thumbnail.png",
                         mime_type="image/png",
                     )
+            elif is_audio:
+                # --- Audio Upload Logic ---
+                # Audio files don't need image processing or aspect ratio
+                final_aspect_ratio = aspect_ratio or AspectRatioEnum.RATIO_1_1
+                
+                # Determine audio mime type
+                audio_mime = file.content_type or "audio/mpeg"
+                file_extension = os.path.splitext(file.filename or "audio.mp3")[1] or ".mp3"
+                
+                # Upload the audio file directly
+                final_gcs_uri = self.gcs_service.store_to_gcs(
+                    folder=f"source_assets/{user.id}/audio",
+                    file_name=f"{file_hash}{file_extension}",
+                    mime_type=audio_mime,
+                    contents=contents,
+                    decode=False,
+                )
             else:
                 # --- Image Upload & Upscale Logic ---
                 # Check for valid aspect ratio early in the process
@@ -334,11 +354,25 @@ async def upload_asset(
                 shutil.rmtree(temp_dir)
 
         # 4. Create and save the new UserAsset document
-        mime_type: MimeTypeEnum = (
-            MimeTypeEnum.VIDEO_MP4
-            if file.content_type == MimeTypeEnum.VIDEO_MP4
-            else MimeTypeEnum.IMAGE_PNG
-        )
+        # Determine mime_type based on content_type
+        content_type = file.content_type or ""
+        if content_type.startswith("video/"):
+            mime_type = MimeTypeEnum.VIDEO_MP4
+        elif content_type.startswith("audio/"):
+            # Map common audio types to enum values
+            if content_type in ["audio/mpeg", "audio/mp3"]:
+                mime_type = MimeTypeEnum.AUDIO_MPEG
+            elif content_type == "audio/wav":
+                mime_type = MimeTypeEnum.AUDIO_WAV
+            elif content_type == "audio/ogg":
+                mime_type = MimeTypeEnum.AUDIO_OGG
+            elif content_type == "audio/webm":
+                mime_type = MimeTypeEnum.AUDIO_WEBM
+            else:
+                # Default to MPEG for unknown audio types
+                mime_type = MimeTypeEnum.AUDIO_MPEG
+        else:
+            mime_type = MimeTypeEnum.IMAGE_PNG
 
         is_admin = UserRoleEnum.ADMIN in user.roles
         final_scope = AssetScopeEnum.PRIVATE