types and flexible naming

App/Editor/Schema.py

@@ -2,8 +2,13 @@ from typing import List, Optional
 from pydantic import BaseModel, HttpUrl
 
 
+class LinkInfo(BaseModel):
+    file_name: str
+    link: HttpUrl
+
+
 class EditorRequest(BaseModel):
-    links: Optional[List[HttpUrl]]
+    links: Optional[List[LinkInfo]]  # List of LinkInfo objects
     script: str
 
 

App/Worker.py

@@ -6,7 +6,7 @@ import time
 import requests
 from App import celery_config, bot
 from typing import List
-from App.Editor.Schema import EditorRequest
+from App.Editor.Schema import EditorRequest, LinkInfo
 from celery.signals import worker_process_init
 
 
@@ -54,24 +54,12 @@ def install_dependencies(directory: str):
 
 
 @celery.task
-def download_assets(links: List[str], temp_dir: str):
-    for i, link in enumerate(links):
-        # Make a request to the server to get the filename and format
-        response = requests.head(link)
-
-        # Extract filename and format from the Content-Disposition header, if available
-        content_disposition = response.headers.get("Content-Disposition")
-        if content_disposition and "filename" in content_disposition:
-            _, params = cgi.parse_header(content_disposition)
-            filename = params["filename"]
-        else:
-            # If Content-Disposition is not available, use the last part of the URL as the filename
-            filename = os.path.basename(urlparse(link).path)
-        public_dir = f"{temp_dir}/public"
-        print(public_dir)
-        # Use the extracted filename to save the file
-
-        download_with_wget(link, public_dir, filename)
+def download_assets(links: List[LinkInfo], temp_dir: str):
+    public_dir = f"{temp_dir}/public"
+    for link in links:
+        file_link = link.link
+        file_name = link.file_name
+        download_with_wget(file_link, public_dir, file_name)
 
 
 @celery.task
@@ -92,7 +80,7 @@ def cleanup_temp_directory(
 
 
 @celery.task
-def celery_task(links, script=""):
+def celery_task(video_task: EditorRequest):
     remotion_app_dir = os.path.join("/srv", "Remotion-app")
     project_id = str(uuid.uuid4())
     temp_dir = f"/tmp/{project_id}"
@@ -101,7 +89,7 @@ def celery_task(links, script=""):
     chain(
         copy_remotion_app.si(remotion_app_dir, temp_dir),
         # install_dependencies.si(temp_dir),
-        download_assets.si(links, temp_dir) if links else None,
+        download_assets.si(video_task.links, temp_dir) if video_task.links else None,
         render_video.si(temp_dir, output_dir),
         cleanup_temp_directory.si(temp_dir, output_dir),
     ).apply_async(