fix youtube links in doc extract

GooeyAI · Feb 13, 2024 · 709331d · 709331d
1 parent eee1578
commit 709331d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/daras_ai/image_input.py b/daras_ai/image_input.py
@@ -112,7 +112,9 @@ def safe_filename(filename: str) -> str:
     return out
 
 
-def truncate_filename(text: str, maxlen: int = 100, sep: str = "...") -> str:
+def truncate_filename(
+    text: str | bytes, maxlen: int = 100, sep: str | bytes = "..."
+) -> str | bytes:
     if len(text) <= maxlen:
         return text
     assert len(sep) <= maxlen

diff --git a/recipes/DocExtract.py b/recipes/DocExtract.py
@@ -390,9 +390,11 @@ def process_source(
     )
 
     content_url = existing_values[Columns.content_url.value]
-    is_pdf = "application/pdf" in doc_meta.mime_type
     is_yt = is_yt_url(webpage_url)
-    is_video = "video/" in doc_meta.mime_type or "audio/" in doc_meta.mime_type
+    is_pdf = doc_meta and "application/pdf" in doc_meta.mime_type
+    is_video = doc_meta and (
+        "video/" in doc_meta.mime_type or "audio/" in doc_meta.mime_type
+    )
     if not content_url:
         yield "Downloading"
         if is_yt: