check file size first and clarify processing logic (#4985)

* check file size first and clarify processing logic

* basic gdrive extraction clariy

* typo

---------

Co-authored-by: Wenxi Onyx <wenxi-onyx@Wenxis-MacBook-Pro.local>

Author: wenxi-onyx

backend/onyx/connectors/google_drive/doc_conversion.py

@@ -129,9 +129,34 @@ def _download_and_extract_sections_basic(
     mime_type = file["mimeType"]
     link = file.get(WEB_VIEW_LINK_KEY, "")
 
-    # skip images if not explicitly enabled
-    if not allow_images and is_gdrive_image_mime_type(mime_type):
-        return []
+    # For non-Google files, download the file
+    # Use the correct API call for downloading files
+    # lazy evaluation to only download the file if necessary
+    def response_call() -> bytes:
+        return download_request(service, file_id)
+
+    if is_gdrive_image_mime_type(mime_type):
+        # Skip images if not explicitly enabled
+        if not allow_images:
+            return []
+
+        # Store images for later processing
+        sections: list[TextSection | ImageSection] = []
+        try:
+            with get_session_with_current_tenant() as db_session:
+                section, embedded_id = store_image_and_create_section(
+                    db_session=db_session,
+                    image_data=response_call(),
+                    file_id=file_id,
+                    display_name=file_name,
+                    media_type=mime_type,
+                    file_origin=FileOrigin.CONNECTOR,
+                    link=link,
+                )
+                sections.append(section)
+        except Exception as e:
+            logger.error(f"Failed to process image {file_name}: {e}")
+        return sections
 
     # For Google Docs, Sheets, and Slides, export as plain text
     if mime_type in GOOGLE_MIME_TYPES_TO_EXPORT:
@@ -154,12 +179,6 @@ def _download_and_extract_sections_basic(
         text = response.decode("utf-8")
         return [TextSection(link=link, text=text)]
 
-    # For other file types, download the file
-    # Use the correct API call for downloading files
-    # lazy evaluation to only download the file if necessary
-    def response_call() -> bytes:
-        return download_request(service, file_id)
-
     # Process based on mime type
     if mime_type == "text/plain":
         try:
@@ -189,25 +208,6 @@ def response_call() -> bytes:
         text = pptx_to_text(io.BytesIO(response_call()), file_name=file_name)
         return [TextSection(link=link, text=text)] if text else []
 
-    elif is_gdrive_image_mime_type(mime_type):
-        # For images, store them for later processing
-        sections: list[TextSection | ImageSection] = []
-        try:
-            with get_session_with_current_tenant() as db_session:
-                section, embedded_id = store_image_and_create_section(
-                    db_session=db_session,
-                    image_data=response_call(),
-                    file_id=file_id,
-                    display_name=file_name,
-                    media_type=mime_type,
-                    file_origin=FileOrigin.CONNECTOR,
-                    link=link,
-                )
-                sections.append(section)
-        except Exception as e:
-            logger.error(f"Failed to process image {file_name}: {e}")
-        return sections
-
     elif mime_type == "application/pdf":
         text, _pdf_meta, images = read_pdf_file(io.BytesIO(response_call()))
         pdf_sections: list[TextSection | ImageSection] = [
@@ -230,27 +230,18 @@ def response_call() -> bytes:
             logger.error(f"Failed to process PDF images in {file_name}: {e}")
         return pdf_sections
 
-    else:
-        # For unsupported file types, try to extract text
-        if mime_type in [
-            "application/vnd.google-apps.video",
-            "application/vnd.google-apps.audio",
-            "application/zip",
-        ]:
-            return []
+    # Final attempt at extracting text
+    file_ext = get_file_ext(file.get("name", ""))
+    if file_ext not in ALL_ACCEPTED_FILE_EXTENSIONS:
+        logger.warning(f"Skipping file {file.get('name')} due to extension.")
+        return []
 
-        # don't download the file at all if it's an unhandled extension
-        file_ext = get_file_ext(file.get("name", ""))
-        if file_ext not in ALL_ACCEPTED_FILE_EXTENSIONS:
-            logger.warning(f"Skipping file {file.get('name')} due to extension.")
-            return []
-        # For unsupported file types, try to extract text
-        try:
-            text = extract_file_text(io.BytesIO(response_call()), file_name)
-            return [TextSection(link=link, text=text)]
-        except Exception as e:
-            logger.warning(f"Failed to extract text from {file_name}: {e}")
-            return []
+    try:
+        text = extract_file_text(io.BytesIO(response_call()), file_name)
+        return [TextSection(link=link, text=text)]
+    except Exception as e:
+        logger.warning(f"Failed to extract text from {file_name}: {e}")
+        return []
 
 
 def _find_nth(haystack: str, needle: str, n: int, start: int = 0) -> int:
@@ -451,6 +442,19 @@ def _get_docs_service() -> GoogleDocsService:
             logger.info("Skipping shortcut/folder.")
             return None
 
+        size_str = file.get("size")
+        if size_str:
+            try:
+                size_int = int(size_str)
+            except ValueError:
+                logger.warning(f"Parsing string to int failed: size_str={size_str}")
+            else:
+                if size_int > size_threshold:
+                    logger.warning(
+                        f"{file.get('name')} exceeds size threshold of {size_threshold}. Skipping."
+                    )
+                    return None
+
         # If it's a Google Doc, we might do advanced parsing
         if file.get("mimeType") == GDriveMimeType.DOC.value:
             try:
@@ -476,22 +480,8 @@ def _get_docs_service() -> GoogleDocsService:
                 logger.warning(
                     f"Error in advanced parsing: {e}. Falling back to basic extraction."
                 )
-
-        size_str = file.get("size")
-        if size_str:
-            try:
-                size_int = int(size_str)
-            except ValueError:
-                logger.warning(f"Parsing string to int failed: size_str={size_str}")
-            else:
-                if size_int > size_threshold:
-                    logger.warning(
-                        f"{file.get('name')} exceeds size threshold of {size_threshold}. Skipping."
-                    )
-                    return None
-
-        # If we don't have sections yet, use the basic extraction method
-        if not sections:
+        # Not Google Doc, attempt basic extraction
+        else:
             sections = _download_and_extract_sections_basic(
                 file, _get_drive_service(), allow_images
             )