diff --git a/backend/onyx/connectors/sharepoint/connector.py b/backend/onyx/connectors/sharepoint/connector.py index 85de5ca058..e9f6c7730f 100644 --- a/backend/onyx/connectors/sharepoint/connector.py +++ b/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,23 +394,26 @@ 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 + 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_{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 + 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) ) - - for idx, (img_data, img_name) in enumerate(extraction_result.embedded_images): - 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_origin=FileOrigin.CONNECTOR, - ) - image_section.link = driveitem.web_url - sections.append(image_section) + # 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}") @@ -729,6 +756,7 @@ class SharepointConnector( 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 @@ class SharepointConnector( 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 @@ class SharepointConnector( 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( diff --git a/backend/onyx/file_processing/extract_file_text.py b/backend/onyx/file_processing/extract_file_text.py index 1b65f5b19c..53efb8545b 100644 --- a/backend/onyx/file_processing/extract_file_text.py +++ b/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