From e968a3d9f8d9af01f0e8d93fffb3490e4d91c959 Mon Sep 17 00:00:00 2001
From: Wenjie Zhang <xerrors@163.com>
Date: Sun, 29 Sep 2024 16:13:06 +0800
Subject: [PATCH] fix pdf parser bug

---
 requirements.txt           |  2 ++
 src/core/database.py       | 10 ++++++++--
 src/core/indexing.py       |  2 +-
 src/plugins/pdf2txt.py     | 29 ++++++++++++++++++++++++++++-
 src/views/database_view.py |  2 +-
 5 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 23d2e274..2def7042 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,6 +7,7 @@ Flask_Cors>=4.0.1
 llama_index>=0.11.8
 openai>=1.44.1
 paddleocr>=2.8.1
+paddlepaddle>=2.4.0
 pandas>=2.2.2
 Pillow>=10.4.0
 pymilvus>=2.4.4
@@ -19,4 +20,5 @@ zhipuai>=2.1.2
 neo4j>=5.22.0
 sentencepiece==0.1.99
 llama-index-readers-file
+opencv-python-headless
 docx2txt
\ No newline at end of file
diff --git a/src/core/database.py b/src/core/database.py
index 59108a44..629b5912 100644
--- a/src/core/database.py
+++ b/src/core/database.py
@@ -1,6 +1,7 @@
 import os
 import json
 import time
+from src.plugins import pdf2txt
 from src.utils import hashstr, setup_logger, is_text_pdf
 from src.models.embedding import get_embedding_model
 
@@ -108,7 +109,7 @@ class DataBaseManager:
                 "file_id": "file_" + hashstr(file + str(time.time())),
                 "filename": os.path.basename(file),
                 "path": file,
-                "type": file.split(".")[-1],
+                "type": file.split(".")[-1].lower(),
                 "status": "waiting",
                 "created_at": time.time()
             }
@@ -122,7 +123,12 @@ class DataBaseManager:
             db.files[idx]["status"] = "processing"
 
             try:
-                nodes = chunk(new_file["path"], params=params)
+                if new_file["type"] == "pdf":
+                    texts = self.read_text(new_file["path"])
+                    nodes = chunk(texts, params=params)
+                else:
+                    nodes = chunk(new_file["path"], params=params)
+
                 self.knowledge_base.add_documents(
                     docs=[node.text for node in nodes],
                     collection_name=db.metaname,
diff --git a/src/core/indexing.py b/src/core/indexing.py
index 36b311b5..a20e480d 100644
--- a/src/core/indexing.py
+++ b/src/core/indexing.py
@@ -18,7 +18,7 @@ def chunk(text_or_path, params=None):
 
     if os.path.isfile(text_or_path) and "uploads" in text_or_path:
         parser = SimpleFileNodeParser()
-        file_type = Path(text_or_path).suffix
+        file_type = Path(text_or_path).suffix.lower()
         if file_type in [".txt", ".json", ".md"]:
             docs = FlatReader().load_data(Path(text_or_path))
         elif file_type in [".docx"]:
diff --git a/src/plugins/pdf2txt.py b/src/plugins/pdf2txt.py
index bb9c4d07..7b70e155 100644
--- a/src/plugins/pdf2txt.py
+++ b/src/plugins/pdf2txt.py
@@ -4,8 +4,9 @@ import fitz  # fitz就是pip install PyMuPDF
 from PIL import Image
 from copy import deepcopy
 from tqdm import tqdm
+from pathlib import Path
 from argparse import ArgumentParser
-from src.utils import logger
+from src.utils import logger, is_text_pdf
 import numpy as np  # Added import for numpy
 
 GOLBAL_STATE = {}
@@ -15,6 +16,9 @@ def pdf2txt(pdf_path, return_text=False):
     if not os.path.exists(pdf_path):
         raise FileNotFoundError(f"File not found: {pdf_path}")
 
+    if is_text_pdf(pdf_path):
+        return pdfreader(pdf_path)
+
     # Importing these modules here to avoid unnecessary imports in other files
     from paddleocr import PPStructure, save_structure_res
     from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx
@@ -108,6 +112,29 @@ def convert_imgs(pdf_path, output_dir):
 def get_state(task_id):
     return GOLBAL_STATE.get(task_id, {})
 
+
+def pdfreader(file_path):
+    """读取PDF文件并返回text文本"""
+    assert os.path.exists(file_path), "File not found"
+    assert file_path.endswith(".pdf"), "File format not supported"
+
+    from llama_index.readers.file import PDFReader
+    doc = PDFReader().load_data(file=Path(file_path))
+
+    # 简单的拼接起来之后返回纯文本
+    text = "\n\n".join([d.get_content() for d in doc])
+    return text
+
+def plainreader(file_path):
+    """读取普通文本文件并返回text文本"""
+    assert os.path.exists(file_path), "File not found"
+
+    with open(file_path, "r") as f:
+        text = f.read()
+    return text
+
+
+
 if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument('--pdf-path', type=str, required=True, help='Path to the PDF file')
diff --git a/src/views/database_view.py b/src/views/database_view.py
index ebaa255b..a904c294 100644
--- a/src/views/database_view.py
+++ b/src/views/database_view.py
@@ -118,7 +118,7 @@ def upload_file():
     if file:
         upload_dir = os.path.join(startup.config.save_dir, "data/uploads")
         os.makedirs(upload_dir, exist_ok=True)
-        filename = f"{hashstr(file.filename, 4, with_salt=True)}_{file.filename}"
+        filename = f"{hashstr(file.filename, 4, with_salt=True)}_{file.filename}".lower()
         file_path = os.path.join(upload_dir, filename)
         file.save(file_path)
         return jsonify({'message': 'File successfully uploaded', 'file_path': file_path}), 200