Mbonea commited on
Commit
7dbba0a
1 Parent(s): 07dca11
Files changed (2) hide show
  1. App/Worker.py +16 -2
  2. requirements.txt +1 -0
App/Worker.py CHANGED
@@ -2,6 +2,7 @@ from celery import Celery, chain
2
  import os, shutil, subprocess
3
  import uuid
4
  import time
 
5
  from App import celery_config, bot
6
  from typing import List
7
  from App.Editor.Schema import EditorRequest
@@ -43,8 +44,21 @@ def install_dependencies(directory: str):
43
  @celery.task
44
  def download_assets(links: List[str], temp_dir: str):
45
  for i, link in enumerate(links):
46
- download_dir = os.path.join(temp_dir, "public")
47
- download_with_wget(link, download_dir)
 
 
 
 
 
 
 
 
 
 
 
 
 
48
 
49
 
50
  @celery.task
 
2
  import os, shutil, subprocess
3
  import uuid
4
  import time
5
+ import requests
6
  from App import celery_config, bot
7
  from typing import List
8
  from App.Editor.Schema import EditorRequest
 
44
  @celery.task
45
  def download_assets(links: List[str], temp_dir: str):
46
  for i, link in enumerate(links):
47
+ # Make a request to the server to get the filename and format
48
+ response = requests.head(link)
49
+
50
+ # Extract filename and format from the Content-Disposition header, if available
51
+ content_disposition = response.headers.get("Content-Disposition")
52
+ if content_disposition and "filename" in content_disposition:
53
+ _, params = cgi.parse_header(content_disposition)
54
+ filename = params["filename"]
55
+ else:
56
+ # If Content-Disposition is not available, use the last part of the URL as the filename
57
+ filename = os.path.basename(urlparse(link).path)
58
+
59
+ # Use the extracted filename to save the file
60
+ destination_path = os.path.join(temp_dir, filename)
61
+ download_with_wget(link, destination_path)
62
 
63
 
64
  @celery.task
requirements.txt CHANGED
@@ -15,6 +15,7 @@ Werkzeug==2.2.2
15
  pydantic[email]
16
  uvicorn==0.21.1
17
  gunicorn
 
18
  ujson
19
  psutil
20
  celery
 
15
  pydantic[email]
16
  uvicorn==0.21.1
17
  gunicorn
18
+ requests
19
  ujson
20
  psutil
21
  celery