Update main.py

main.py

@@ -4,15 +4,21 @@ import time
 import logging
 from pathlib import Path
 from typing import List, Dict, Optional
-from dataclasses import dataclass, asdict
-
-from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
-from magic_pdf.data.dataset import PymuDocDataset
-from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
-from magic_pdf.config.enums import SupportedPdfParseMethod
+from dataclasses import dataclass
+from fastapi.encoders import jsonable_encoder
 from sentence_transformers import SentenceTransformer
 from llama_cpp import Llama
-from fastapi.encoders import jsonable_encoder
+
+# Fix: Dynamically adjust the module path if magic_pdf is in a non-standard location
+try:
+    from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
+    from magic_pdf.data.dataset import PymuDocDataset
+    from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
+    from magic_pdf.config.enums import SupportedPdfParseMethod
+except ModuleNotFoundError as e:
+    logging.error(f"Failed to import magic_pdf modules: {e}")
+    logging.info("Ensure that the magic_pdf package is installed and accessible in your Python environment.")
+    raise e
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -57,39 +63,43 @@ class PDFProcessor:
         
         os.makedirs(local_image_dir, exist_ok=True)
         
-        image_writer = FileBasedDataWriter(str(local_image_dir))
-        md_writer = FileBasedDataWriter(str(local_md_dir))
-        
-        # Read PDF
-        reader = FileBasedDataReader("")
-        pdf_bytes = reader.read(pdf_path)
-        
-        # Create dataset and process
-        ds = PymuDocDataset(pdf_bytes)
-        
-        if ds.classify() == SupportedPdfParseMethod.OCR:
-            infer_result = ds.apply(doc_analyze, ocr=True)
-            pipe_result = infer_result.pipe_ocr_mode(image_writer)
-        else:
-            infer_result = ds.apply(doc_analyze, ocr=False)
-            pipe_result = infer_result.pipe_txt_mode(image_writer)
-        
-        # Get structured content
-        middle_json = pipe_result.get_middle_json()
-        tables = self._extract_tables(middle_json)
-        text_blocks = self._extract_text_blocks(middle_json)
-        
-        # Process text blocks with LLM
-        products = []
-        for block in text_blocks:
-            product = self._process_text_block(block)
-            if product:
-                product.tables = tables
-                products.append(product.to_dict())
-        
-        logger.info(f"Processed {len(products)} products in {time.time()-start_time:.2f}s")
-        return {"products": products, "tables": tables}
-
+        try:
+            image_writer = FileBasedDataWriter(str(local_image_dir))
+            md_writer = FileBasedDataWriter(str(local_md_dir))
+            
+            # Read PDF
+            reader = FileBasedDataReader("")
+            pdf_bytes = reader.read(pdf_path)
+            
+            # Create dataset and process
+            ds = PymuDocDataset(pdf_bytes)
+            
+            if ds.classify() == SupportedPdfParseMethod.OCR:
+                infer_result = ds.apply(doc_analyze, ocr=True)
+                pipe_result = infer_result.pipe_ocr_mode(image_writer)
+            else:
+                infer_result = ds.apply(doc_analyze, ocr=False)
+                pipe_result = infer_result.pipe_txt_mode(image_writer)
+            
+            # Get structured content
+            middle_json = pipe_result.get_middle_json()
+            tables = self._extract_tables(middle_json)
+            text_blocks = self._extract_text_blocks(middle_json)
+            
+            # Process text blocks with LLM
+            products = []
+            for block in text_blocks:
+                product = self._process_text_block(block)
+                if product:
+                    product.tables = tables
+                    products.append(product.to_dict())
+            
+            logger.info(f"Processed {len(products)} products in {time.time()-start_time:.2f}s")
+            return {"products": products, "tables": tables}
+        except Exception as e:
+            logger.error(f"Error during PDF processing: {e}")
+            raise RuntimeError("PDF processing failed.") from e
+    
     def _extract_tables(self, middle_json: Dict) -> List[Dict]:
         """Extract tables from MinerU's middle JSON"""
         tables = []
@@ -102,7 +112,7 @@ class PDFProcessor:
                     "content": table.get('content', [])
                 })
         return tables
-
+    
     def _extract_text_blocks(self, middle_json: Dict) -> List[str]:
         """Extract text blocks from MinerU's middle JSON"""
         text_blocks = []
@@ -111,7 +121,7 @@ class PDFProcessor:
                 if block.get('type') == 'text':
                     text_blocks.append(block.get('text', ''))
         return text_blocks
-
+    
     def _process_text_block(self, text: str) -> Optional[ProductSpec]:
         """Process text block with LLM"""
         prompt = self._generate_query_prompt(text)
@@ -126,12 +136,11 @@ class PDFProcessor:
         except Exception as e:
             logger.warning(f"Error processing text block: {e}")
             return None
-
+    
     def _generate_query_prompt(self, text: str) -> str:
         """Generate extraction prompt"""
         return f"""Extract product specifications from this text:
 {text}
-
 Return JSON format:
 {{
     "name": "product name",
@@ -139,7 +148,7 @@ Return JSON format:
     "price": numeric_price,
     "attributes": {{ "key": "value" }}
 }}"""
-
+    
     def _parse_response(self, response: str) -> Optional[ProductSpec]:
         """Parse LLM response"""
         try: