refac

backend/open_webui/retrieval/loaders/mineru.py

@@ -28,20 +28,22 @@ class MinerULoader:
         api_key: str = '',
         params: dict = None,
         timeout: Optional[int] = 300,
+        max_markdown_bytes: Optional[int] = None,
     ):
         self.file_path = file_path
         self.api_mode = api_mode.lower()
         self.api_url = api_url.rstrip('/')
         self.api_key = api_key
         self.timeout = timeout
+        self.max_markdown_bytes = max_markdown_bytes
 
         # Parse params dict with defaults
         self.params = params or {}
-        self.enable_ocr = params.get('enable_ocr', False)
+        self.enable_ocr = self.params.get('enable_ocr', False)
-        self.enable_formula = params.get('enable_formula', True)
+        self.enable_formula = self.params.get('enable_formula', True)
-        self.enable_table = params.get('enable_table', True)
+        self.enable_table = self.params.get('enable_table', True)
-        self.language = params.get('language', 'en')
+        self.language = self.params.get('language', 'en')
-        self.model_version = params.get('model_version', 'pipeline')
+        self.model_version = self.params.get('model_version', 'pipeline')
 
         self.page_ranges = self.params.pop('page_ranges', '')
 
@@ -435,67 +437,77 @@ class MinerULoader:
                 detail=f'Error downloading results: {str(e)}',
             )
 
-        # Save ZIP to temporary file and extract
+        # Save ZIP to temporary file before reading.
+        tmp_zip_path = None
+        markdown_content = None
         try:
             with tempfile.NamedTemporaryFile(delete=False, suffix='.zip') as tmp_zip:
                 tmp_zip.write(response.content)
                 tmp_zip_path = tmp_zip.name
 
-            with tempfile.TemporaryDirectory() as tmp_dir:
+            with zipfile.ZipFile(tmp_zip_path, 'r') as zip_ref:
-                # Extract ZIP
+                members = zip_ref.infolist()
-                with zipfile.ZipFile(tmp_zip_path, 'r') as zip_ref:
+                all_files = [member.filename for member in members]
-                    zip_ref.extractall(tmp_dir)
+                md_members = [member for member in members if member.filename.endswith('.md')]
+                read_errors = []
 
-                # Find markdown file - search recursively for any .md file
+                for member in md_members:
-                markdown_content = None
+                    log.info(f'Found markdown file in ZIP: {member.filename}')
-                found_md_path = None
+                    try:
-
+                        with zip_ref.open(member, 'r') as f:
-                # First, list all files in the ZIP for debugging
+                            if self.max_markdown_bytes is None:
-                all_files = []
+                                content = f.read()
-                for root, dirs, files in os.walk(tmp_dir):
+                            else:
-                    for file in files:
+                                content = f.read(self.max_markdown_bytes + 1)
-                        full_path = os.path.join(root, file)
+                                if len(content) > self.max_markdown_bytes:
-                        all_files.append(full_path)
+                                    raise HTTPException(
-                        # Look for any .md file
+                                        status.HTTP_502_BAD_GATEWAY,
-                        if file.endswith('.md'):
+                                        detail=f'Markdown file in results ZIP is too large: {member.filename}',
-                            found_md_path = full_path
+                                    )
-                            log.info(f'Found markdown file at: {full_path}')
+                            markdown_content = content.decode('utf-8')
-                            try:
+                    except UnicodeDecodeError as e:
-                                with open(full_path, 'r', encoding='utf-8') as f:
+                        read_errors.append(f'{member.filename}: {e}')
-                                    markdown_content = f.read()
+                        log.warning(f'Failed to decode {member.filename}: {e}')
-                                if markdown_content:  # Use the first non-empty markdown file
+                        continue
-                                    break
+                    except HTTPException:
-                            except Exception as e:
+                        raise
-                                log.warning(f'Failed to read {full_path}: {e}')
+                    except Exception as e:
+                        read_errors.append(f'{member.filename}: {e}')
+                        log.warning(f'Failed to read {member.filename}: {e}')
+                        continue
                     if markdown_content:
                         break
 
                 if markdown_content is None:
                     log.error(f'Available files in ZIP: {all_files}')
-                    # Try to provide more helpful error message
+                    if read_errors:
-                    md_files = [f for f in all_files if f.endswith('.md')]
+                        error_msg = f"Found .md files but couldn't read them: {read_errors}"
-                    if md_files:
-                        error_msg = f"Found .md files but couldn't read them: {md_files}"
                     else:
                         error_msg = f'No .md files found in ZIP. Available files: {all_files}'
                     raise HTTPException(
                         status.HTTP_502_BAD_GATEWAY,
                         detail=error_msg,
                     )
-
-            # Clean up temporary ZIP file
-            os.unlink(tmp_zip_path)
-
         except zipfile.BadZipFile as e:
             raise HTTPException(
                 status.HTTP_502_BAD_GATEWAY,
                 detail=f'Invalid ZIP file received: {e}',
             )
+        except HTTPException:
+            raise
         except Exception as e:
             raise HTTPException(
                 status.HTTP_500_INTERNAL_SERVER_ERROR,
                 detail=f'Error extracting ZIP: {str(e)}',
             )
+        finally:
+            if tmp_zip_path:
+                try:
+                    os.unlink(tmp_zip_path)
+                except FileNotFoundError:
+                    pass
+                except Exception as e:
+                    log.warning(f'Failed to remove temporary ZIP file {tmp_zip_path}: {e}')
 
         if not markdown_content:
             raise HTTPException(