fix: skip huge files on sdk fallback (#5421)

Author: evan-onyx

backend/onyx/connectors/sharepoint/connector.py

@@ -142,6 +142,10 @@ class SharepointAuthMethod(Enum):
     CERTIFICATE = "certificate"
 
 
+class SizeCapExceeded(Exception):
+    """Exception raised when the size cap is exceeded."""
+
+
 def load_certificate_from_pfx(pfx_data: bytes, password: str) -> CertificateData | None:
     """Load certificate from .pfx file for MSAL authentication"""
     try:
@@ -240,7 +244,7 @@ def _download_with_cap(url: str, timeout: int, cap: int) -> bytes:
     Behavior:
     - Checks `Content-Length` first and aborts early if it exceeds `cap`.
     - Otherwise streams the body in chunks and stops once `cap` is surpassed.
-    - Raises `RuntimeError('size_cap_exceeded')` when the cap would be exceeded.
+    - Raises `SizeCapExceeded` when the cap would be exceeded.
     - Returns the full bytes if the content fits within `cap`.
     """
     with requests.get(url, stream=True, timeout=timeout) as resp:
@@ -254,7 +258,7 @@ def _download_with_cap(url: str, timeout: int, cap: int) -> bytes:
                 logger.warning(
                     f"Content-Length {content_len} exceeds cap {cap}; skipping download."
                 )
-                raise RuntimeError("size_cap_exceeded")
+                raise SizeCapExceeded("pre_download")
 
         buf = io.BytesIO()
         # Stream in 64KB chunks; adjust if needed for slower networks.
@@ -267,11 +271,32 @@ def _download_with_cap(url: str, timeout: int, cap: int) -> bytes:
                 logger.warning(
                     f"Streaming download exceeded cap {cap} bytes; aborting early."
                 )
-                raise RuntimeError("size_cap_exceeded")
+                raise SizeCapExceeded("during_download")
 
         return buf.getvalue()
 
 
+def _download_via_sdk_with_cap(
+    driveitem: DriveItem, bytes_allowed: int, chunk_size: int = 64 * 1024
+) -> bytes:
+    """Use the Office365 SDK streaming download with a hard byte cap.
+
+    Raises SizeCapExceeded("during_sdk_download") if the cap would be exceeded.
+    """
+    buf = io.BytesIO()
+
+    def on_chunk(bytes_read: int) -> None:
+        # bytes_read is total bytes seen so far per SDK contract
+        if bytes_read > bytes_allowed:
+            raise SizeCapExceeded("during_sdk_download")
+
+    # modifies the driveitem to change its download behavior
+    driveitem.download_session(buf, chunk_downloaded=on_chunk, chunk_size=chunk_size)
+    # Execute the configured request with retries using existing helper
+    sleep_and_retry(driveitem.context, "download_session")
+    return buf.getvalue()
+
+
 def _convert_driveitem_to_document_with_permissions(
     driveitem: DriveItem,
     drive_name: str,
@@ -322,19 +347,16 @@ def _convert_driveitem_to_document_with_permissions(
     content_bytes: bytes | None = None
     if download_url:
         try:
+            # Use this to test the sdk size cap
+            # raise requests.RequestException("test")
             content_bytes = _download_with_cap(
                 download_url,
                 REQUEST_TIMEOUT_SECONDS,
                 SHAREPOINT_CONNECTOR_SIZE_THRESHOLD,
             )
-        except RuntimeError as e:
-            if "size_cap_exceeded" in str(e):
-                logger.warning(
-                    f"Skipping '{driveitem.name}' exceeded size cap during streaming."
-                )
-                return None
-            else:
-                raise
+        except SizeCapExceeded as e:
+            logger.warning(f"Skipping '{driveitem.name}' exceeded size cap: {str(e)}")
+            return None
         except requests.RequestException as e:
             status = e.response.status_code if e.response is not None else -1
             logger.warning(
@@ -343,13 +365,15 @@ def _convert_driveitem_to_document_with_permissions(
 
     # Fallback to SDK content if needed
     if content_bytes is None:
-        content = sleep_and_retry(driveitem.get_content(), "get_content")
-        if content is None or not isinstance(
-            getattr(content, "value", None), (bytes, bytearray)
-        ):
-            logger.warning(f"Could not access content for '{driveitem.name}'")
-            raise ValueError(f"Could not access content for '{driveitem.name}'")
-        content_bytes = bytes(content.value)
+        try:
+            content_bytes = _download_via_sdk_with_cap(
+                driveitem, SHAREPOINT_CONNECTOR_SIZE_THRESHOLD
+            )
+        except SizeCapExceeded:
+            logger.warning(
+                f"Skipping '{driveitem.name}' exceeded size cap during SDK streaming."
+            )
+            return None
 
     sections: list[TextSection | ImageSection] = []
     file_ext = driveitem.name.split(".")[-1]
@@ -370,24 +394,27 @@ def _convert_driveitem_to_document_with_permissions(
         sections.append(image_section)
     else:
         # Note: we don't process Onyx metadata for connectors like Drive & Sharepoint, but could
-        extraction_result = extract_text_and_images(
-            file=io.BytesIO(content_bytes), file_name=driveitem.name
-        )
-        if extraction_result.text_content:
-            sections.append(
-                TextSection(link=driveitem.web_url, text=extraction_result.text_content)
-            )
-
-        for idx, (img_data, img_name) in enumerate(extraction_result.embedded_images):
+        def _store_embedded_image(img_data: bytes, img_name: str) -> None:
             image_section, _ = store_image_and_create_section(
                 image_data=img_data,
-                file_id=f"{driveitem.id}_img_{idx}",
-                display_name=img_name or f"{driveitem.name} - image {idx}",
+                file_id=f"{driveitem.id}_img_{len(sections)}",
+                display_name=img_name or f"{driveitem.name} - image {len(sections)}",
                 file_origin=FileOrigin.CONNECTOR,
             )
             image_section.link = driveitem.web_url
             sections.append(image_section)
 
+        extraction_result = extract_text_and_images(
+            file=io.BytesIO(content_bytes),
+            file_name=driveitem.name,
+            image_callback=_store_embedded_image,
+        )
+        if extraction_result.text_content:
+            sections.append(
+                TextSection(link=driveitem.web_url, text=extraction_result.text_content)
+            )
+        # Any embedded images were stored via the callback; the returned list may be empty.
+
     if include_permissions and ctx is not None:
         logger.info(f"Getting external access for {driveitem.name}")
         external_access = get_sharepoint_external_access(
@@ -729,6 +756,7 @@ def _get_drive_items_for_drive_name(
                     for folder_part in site_descriptor.folder_path.split("/"):
                         root_folder = root_folder.get_by_path(folder_part)
 
+                # TODO: consider ways to avoid materializing the entire list of files in memory
                 query = root_folder.get_files(
                     recursive=True,
                     page_size=1000,
@@ -837,6 +865,7 @@ def _fetch_driveitems(
                             root_folder = root_folder.get_by_path(folder_part)
 
                     # Get all items recursively
+                    # TODO: consider ways to avoid materializing the entire list of files in memory
                     query = root_folder.get_files(
                         recursive=True,
                         page_size=1000,
@@ -985,6 +1014,8 @@ def _fetch_site_pages(
         all_pages = pages_data.get("value", [])
 
         # Handle pagination if there are more pages
+        # TODO: This accumulates all pages in memory and can be heavy on large tenants.
+        #       We should process each page incrementally to avoid unbounded growth.
         while "@odata.nextLink" in pages_data:
             next_url = pages_data["@odata.nextLink"]
             response = requests.get(

backend/onyx/file_processing/extract_file_text.py

@@ -242,7 +242,10 @@ def pdf_to_text(file: IO[Any], pdf_pass: str | None = None) -> str:
 
 
 def read_pdf_file(
-    file: IO[Any], pdf_pass: str | None = None, extract_images: bool = False
+    file: IO[Any],
+    pdf_pass: str | None = None,
+    extract_images: bool = False,
+    image_callback: Callable[[bytes, str], None] | None = None,
 ) -> tuple[str, dict[str, Any], Sequence[tuple[bytes, str]]]:
     """
     Returns the text, basic PDF metadata, and optionally extracted images.
@@ -292,7 +295,11 @@ def read_pdf_file(
                         f"page_{page_num + 1}_image_{image_file_object.name}."
                         f"{image.format.lower() if image.format else 'png'}"
                     )
-                    extracted_images.append((img_bytes, image_name))
+                    if image_callback is not None:
+                        # Stream image out immediately
+                        image_callback(img_bytes, image_name)
+                    else:
+                        extracted_images.append((img_bytes, image_name))
 
         return text, metadata, extracted_images
 
@@ -304,28 +311,32 @@ def read_pdf_file(
     return "", metadata, []
 
 
-def extract_docx_images(docx_bytes: IO[Any]) -> list[tuple[bytes, str]]:
+def extract_docx_images(docx_bytes: IO[Any]) -> Iterator[tuple[bytes, str]]:
     """
     Given the bytes of a docx file, extract all the images.
     Returns a list of tuples (image_bytes, image_name).
     """
-    out = []
     try:
         with zipfile.ZipFile(docx_bytes) as z:
             for name in z.namelist():
                 if name.startswith("word/media/"):
-                    out.append((z.read(name), name.split("/")[-1]))
+                    yield (z.read(name), name.split("/")[-1])
     except Exception:
         logger.exception("Failed to extract all docx images")
-    return out
 
 
 def docx_to_text_and_images(
-    file: IO[Any], file_name: str = ""
+    file: IO[Any],
+    file_name: str = "",
+    image_callback: Callable[[bytes, str], None] | None = None,
 ) -> tuple[str, Sequence[tuple[bytes, str]]]:
     """
     Extract text from a docx.
     Return (text_content, list_of_images).
+
+    The caller can choose to provide a callback to handle images with the intent
+    of avoiding materializing the list of images in memory.
+    The images list returned is empty in this case.
     """
     md = MarkItDown(enable_plugins=False)
     try:
@@ -349,7 +360,15 @@ def docx_to_text_and_images(
         return text_content_raw or "", []
 
     file.seek(0)
-    return doc.markdown, extract_docx_images(to_bytesio(file))
+    if image_callback is None:
+        return doc.markdown, list(extract_docx_images(to_bytesio(file)))
+    # If a callback is provided, iterate and stream images without accumulating
+    try:
+        for img_file_bytes, img_file_name in extract_docx_images(to_bytesio(file)):
+            image_callback(img_file_bytes, img_file_name)
+    except Exception:
+        logger.exception("Failed to stream docx images")
+    return doc.markdown, []
 
 
 def pptx_to_text(file: IO[Any], file_name: str = "") -> str:
@@ -506,6 +525,7 @@ def extract_text_and_images(
     file_name: str,
     pdf_pass: str | None = None,
     content_type: str | None = None,
+    image_callback: Callable[[bytes, str], None] | None = None,
 ) -> ExtractionResult:
     """
     Primary new function for the updated connector.
@@ -539,7 +559,9 @@ def extract_text_and_images(
 
         # docx example for embedded images
         if extension == ".docx":
-            text_content, images = docx_to_text_and_images(file, file_name)
+            text_content, images = docx_to_text_and_images(
+                file, file_name, image_callback=image_callback
+            )
             return ExtractionResult(
                 text_content=text_content, embedded_images=images, metadata={}
             )
@@ -551,6 +573,7 @@ def extract_text_and_images(
                 file,
                 pdf_pass,
                 extract_images=get_image_extraction_and_analysis_enabled(),
+                image_callback=image_callback,
             )
             return ExtractionResult(
                 text_content=text_content, embedded_images=images, metadata=pdf_metadata