Update modules/knowledge_base/processor.py

modules/knowledge_base/processor.py

@@ -1,230 +1,230 @@
-from typing import List, Dict, Callable, Optional
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.document_loaders import (
-    DirectoryLoader,
-    UnstructuredMarkdownLoader,
-    PyPDFLoader,
-    TextLoader
-)
-import os
-import requests
-import base64
-from PIL import Image
-import io
-
-class DocumentLoader:
-    """通用文档加载器"""
-    def __init__(self, file_path: str, original_filename: str = None):
-        self.file_path = file_path
-        # 使用传入的原始文件名或者从路径提取的文件名
-        self.original_filename = original_filename or os.path.basename(file_path)
-        # 从原始文件名中获取扩展名，确保中文文件名也能正确识别文件类型
-        self.extension = os.path.splitext(self.original_filename)[1].lower()
-        self.api_key = os.getenv("API_KEY")
-        self.api_base = os.getenv("BASE_URL")
-        
-    def process_image(self, image_path: str) -> str:
-        """使用 SiliconFlow VLM 模型处理图片"""
-        try:
-            # 读取图片并转换为base64
-            with open(image_path, 'rb') as image_file:
-                image_data = image_file.read()
-                base64_image = base64.b64encode(image_data).decode('utf-8')
-            
-            # 调用 SiliconFlow API
-            headers = {
-                "Authorization": f"Bearer {self.api_key}",
-                "Content-Type": "application/json"
-            }
-            
-            response = requests.post(
-                f"{self.api_base}/chat/completions",
-                headers=headers,
-                json={
-                    "model": "Qwen/Qwen2.5-VL-72B-Instruct",
-                    "messages": [
-                        {
-                            "role": "user",
-                            "content": [
-                                {
-                                    "type": "image_url",
-                                    "image_url": {
-                                        "url": f"data:image/jpeg;base64,{base64_image}",
-                                        "detail": "high"
-                                    }
-                                },
-                                {
-                                    "type": "text",
-                                    "text": "请详细描述这张图片的内容，包括主要对象、场景、活动、颜色、布局等关键信息。"
-                                }
-                            ]
-                        }
-                    ],
-                    "temperature": 0.7,
-                    "max_tokens": 500
-                }
-            )
-            
-            if response.status_code != 200:
-                raise Exception(f"图片处理API调用失败: {response.text}")
-                
-            description = response.json()["choices"][0]["message"]["content"]
-            return description
-            
-        except Exception as e:
-            print(f"处理图片时出错: {str(e)}")
-            return "图片处理失败"
-    
-    def load(self):
-        try:
-            print(f"正在加载文件: {self.file_path}, 原始文件名: {self.original_filename}, 扩展名: {self.extension}")
-            
-            if self.extension == '.md':
-                try:
-                    loader = UnstructuredMarkdownLoader(self.file_path, encoding='utf-8')
-                    return loader.load()
-                except UnicodeDecodeError:
-                    # 如果UTF-8失败，尝试GBK
-                    loader = UnstructuredMarkdownLoader(self.file_path, encoding='gbk')
-                    return loader.load()
-            elif self.extension == '.pdf':
-                loader = PyPDFLoader(self.file_path)
-                return loader.load()
-            elif self.extension == '.txt':
-                try:
-                    loader = TextLoader(self.file_path, encoding='utf-8')
-                    return loader.load()
-                except UnicodeDecodeError:
-                    # 如果UTF-8失败，尝试GBK
-                    loader = TextLoader(self.file_path, encoding='gbk')
-                    return loader.load()
-            elif self.extension in ['.png', '.jpg', '.jpeg', '.gif', '.bmp']:
-                # 处理图片
-                description = self.process_image(self.file_path)
-                # 创建一个包含图片描述的文档
-                from langchain.schema import Document
-                doc = Document(
-                    page_content=description,
-                    metadata={
-                        'source': self.file_path,
-                        'file_name': self.original_filename,  # 使用原始文件名
-                        'img_url': os.path.abspath(self.file_path)  # 存储图片的绝对路径
-                    }
-                )
-                return [doc]
-            else:
-                print(f"不支持的文件扩展名: {self.extension}")
-                raise ValueError(f"不支持的文件格式: {self.extension}")
-                
-        except UnicodeDecodeError:
-            # 如果默认编码处理失败，尝试其他编码
-            print(f"文件编码错误，尝试其他编码: {self.file_path}")
-            if self.extension in ['.md', '.txt']:
-                try:
-                    loader = TextLoader(self.file_path, encoding='gbk')
-                    return loader.load()
-                except Exception as e:
-                    print(f"尝试GBK编码也失败: {str(e)}")
-            raise
-        except Exception as e:
-            print(f"加载文件 {self.file_path} 时出错: {str(e)}")
-            import traceback
-            traceback.print_exc()
-            raise
-
-class DocumentProcessor:
-    def __init__(self):
-        self.text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=1000,
-            chunk_overlap=200,
-            length_function=len,
-        )
-        
-    def get_index_name(self, path: str) -> str:
-        """根据文件路径生成索引名称"""
-        if os.path.isdir(path):
-            # 如果是目录，使用目录名
-            return f"rag_{os.path.basename(path).lower()}"
-        else:
-            # 如果是文件，使用文件名（不含扩展名）
-            return f"rag_{os.path.splitext(os.path.basename(path))[0].lower()}"
-        
-    def process(self, path: str, progress_callback: Optional[Callable] = None, original_filename: str = None) -> List[Dict]:
-        """
-        加载并处理文档，支持目录或单个文件
-        参数：
-            path: 文档路径
-            progress_callback: 进度回调函数，用于报告处理进度
-            original_filename: 原始文件名（包括中文）
-        返回：处理后的文档列表
-        """
-        if os.path.isdir(path):
-            documents = []
-            total_files = sum([len(files) for _, _, files in os.walk(path)])
-            processed_files = 0
-            processed_size = 0
-            
-            for root, _, files in os.walk(path):
-                for file in files:
-                    file_path = os.path.join(root, file)
-                    try:
-                        # 更新处理进度
-                        if progress_callback:
-                            file_size = os.path.getsize(file_path)
-                            processed_size += file_size
-                            processed_files += 1
-                            progress_callback(processed_size, f"处理文件 {processed_files}/{total_files}: {file}")
-                            
-                        # 为目录中的每个文件，传递原始文件名
-                        loader = DocumentLoader(file_path, original_filename=file)
-                        docs = loader.load()
-                        # 添加文件名到metadata
-                        for doc in docs:
-                            doc.metadata['file_name'] = file  # 使用原始文件名
-                        documents.extend(docs)
-                    except Exception as e:
-                        print(f"警告：加载文件 {file_path} 时出错: {str(e)}")
-                        continue
-        else:
-            try:
-                if progress_callback:
-                    file_size = os.path.getsize(path)
-                    progress_callback(file_size * 0.3, f"加载文件: {original_filename or os.path.basename(path)}")
-                    
-                # 为单个文件，传递原始文件名
-                loader = DocumentLoader(path, original_filename=original_filename)
-                documents = loader.load()
-                
-                # 更新进度
-                if progress_callback:
-                    progress_callback(file_size * 0.6, f"处理文件内容...")
-                    
-                # 使用原始文件名而不是存储的文件名
-                file_name = original_filename or os.path.basename(path)
-                for doc in documents:
-                    doc.metadata['file_name'] = file_name
-            except Exception as e:
-                print(f"加载文件时出错: {str(e)}")
-                raise
-        
-        # 分块
-        chunks = self.text_splitter.split_documents(documents)
-        
-        # 更新进度
-        if progress_callback:
-            if os.path.isdir(path):
-                progress_callback(processed_size, f"文档分块完成，共{len(chunks)}个文档片段")
-            else:
-                file_size = os.path.getsize(path)
-                progress_callback(file_size * 0.9, f"文档分块完成，共{len(chunks)}个文档片段")
-        
-        # 处理成统一格式
-        processed_docs = []
-        for i, chunk in enumerate(chunks):
-            processed_docs.append({
-                'id': f'doc_{i}',
-                'content': chunk.page_content,
-                'metadata': chunk.metadata
-            })
-            
+from typing import List, Dict, Callable, Optional
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import (
+    DirectoryLoader,
+    UnstructuredMarkdownLoader,
+    PyPDFLoader,
+    TextLoader
+)
+import os
+import requests
+import base64
+from PIL import Image
+import io
+
+class DocumentLoader:
+    """通用文档加载器"""
+    def __init__(self, file_path: str, original_filename: str = None):
+        self.file_path = file_path
+        # 使用传入的原始文件名或者从路径提取的文件名
+        self.original_filename = original_filename or os.path.basename(file_path)
+        # 从原始文件名中获取扩展名，确保中文文件名也能正确识别文件类型
+        self.extension = os.path.splitext(self.original_filename)[1].lower()
+        self.api_key = os.getenv("API_KEY")
+        self.api_base = os.getenv("BASE_URL")
+        
+    def process_image(self, image_path: str) -> str:
+        """使用 SiliconFlow VLM 模型处理图片"""
+        try:
+            # 读取图片并转换为base64
+            with open(image_path, 'rb') as image_file:
+                image_data = image_file.read()
+                base64_image = base64.b64encode(image_data).decode('utf-8')
+            
+            # 调用 SiliconFlow API
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json"
+            }
+            
+            response = requests.post(
+                f"{self.api_base}/chat/completions",
+                headers=headers,
+                json={
+                    "model": "Qwen/Qwen2.5-VL-72B-Instruct",
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": f"data:image/jpeg;base64,{base64_image}",
+                                        "detail": "high"
+                                    }
+                                },
+                                {
+                                    "type": "text",
+                                    "text": "请详细描述这张图片的内容，包括主要对象、场景、活动、颜色、布局等关键信息。"
+                                }
+                            ]
+                        }
+                    ],
+                    "temperature": 0.7,
+                    "max_tokens": 500
+                }
+            )
+            
+            if response.status_code != 200:
+                raise Exception(f"图片处理API调用失败: {response.text}")
+                
+            description = response.json()["choices"][0]["message"]["content"]
+            return description
+            
+        except Exception as e:
+            print(f"处理图片时出错: {str(e)}")
+            return "图片处理失败"
+    
+    def load(self):
+        try:
+            print(f"正在加载文件: {self.file_path}, 原始文件名: {self.original_filename}, 扩展名: {self.extension}")
+            
+            if self.extension == '.md':
+                try:
+                    loader = UnstructuredMarkdownLoader(self.file_path, encoding='utf-8')
+                    return loader.load()
+                except UnicodeDecodeError:
+                    # 如果UTF-8失败，尝试GBK
+                    loader = UnstructuredMarkdownLoader(self.file_path, encoding='gbk')
+                    return loader.load()
+            elif self.extension == '.pdf':
+                loader = PyPDFLoader(self.file_path)
+                return loader.load()
+            elif self.extension == '.txt':
+                try:
+                    loader = TextLoader(self.file_path, encoding='utf-8')
+                    return loader.load()
+                except UnicodeDecodeError:
+                    # 如果UTF-8失败，尝试GBK
+                    loader = TextLoader(self.file_path, encoding='gbk')
+                    return loader.load()
+            elif self.extension in ['.png', '.jpg', '.jpeg', '.gif', '.bmp']:
+                # 处理图片
+                description = self.process_image(self.file_path)
+                # 创建一个包含图片描述的文档
+                from langchain.schema import Document
+                doc = Document(
+                    page_content=description,
+                    metadata={
+                        'source': self.file_path,
+                        'file_name': self.original_filename,  # 使用原始文件名
+                        'img_url': os.path.abspath(self.file_path)  # 存储图片的绝对路径
+                    }
+                )
+                return [doc]
+            else:
+                print(f"不支持的文件扩展名: {self.extension}")
+                raise ValueError(f"不支持的文件格式: {self.extension}")
+                
+        except UnicodeDecodeError:
+            # 如果默认编码处理失败，尝试其他编码
+            print(f"文件编码错误，尝试其他编码: {self.file_path}")
+            if self.extension in ['.md', '.txt']:
+                try:
+                    loader = TextLoader(self.file_path, encoding='gbk')
+                    return loader.load()
+                except Exception as e:
+                    print(f"尝试GBK编码也失败: {str(e)}")
+            raise
+        except Exception as e:
+            print(f"加载文件 {self.file_path} 时出错: {str(e)}")
+            import traceback
+            traceback.print_exc()
+            raise
+
+class DocumentProcessor:
+    def __init__(self):
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200,
+            length_function=len,
+        )
+        
+    def get_index_name(self, path: str) -> str:
+        """根据文件路径生成索引名称"""
+        if os.path.isdir(path):
+            # 如果是目录，使用目录名
+            return f"rag_{os.path.basename(path).lower()}"
+        else:
+            # 如果是文件，使用文件名（不含扩展名）
+            return f"rag_{os.path.splitext(os.path.basename(path))[0].lower()}"
+        
+    def process(self, path: str, progress_callback: Optional[Callable] = None, original_filename: str = None) -> List[Dict]:
+        """
+        加载并处理文档，支持目录或单个文件
+        参数：
+            path: 文档路径
+            progress_callback: 进度回调函数，用于报告处理进度
+            original_filename: 原始文件名（包括中文）
+        返回：处理后的文档列表
+        """
+        if os.path.isdir(path):
+            documents = []
+            total_files = sum([len(files) for _, _, files in os.walk(path)])
+            processed_files = 0
+            processed_size = 0
+            
+            for root, _, files in os.walk(path):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    try:
+                        # 更新处理进度
+                        if progress_callback:
+                            file_size = os.path.getsize(file_path)
+                            processed_size += file_size
+                            processed_files += 1
+                            progress_callback(processed_size, f"处理文件 {processed_files}/{total_files}: {file}")
+                            
+                        # 为目录中的每个文件，传递原始文件名
+                        loader = DocumentLoader(file_path, original_filename=file)
+                        docs = loader.load()
+                        # 添加文件名到metadata
+                        for doc in docs:
+                            doc.metadata['file_name'] = file  # 使用原始文件名
+                        documents.extend(docs)
+                    except Exception as e:
+                        print(f"警告：加载文件 {file_path} 时出错: {str(e)}")
+                        continue
+        else:
+            try:
+                if progress_callback:
+                    file_size = os.path.getsize(path)
+                    progress_callback(file_size * 0.3, f"加载文件: {original_filename or os.path.basename(path)}")
+                    
+                # 为单个文件，传递原始文件名
+                loader = DocumentLoader(path, original_filename=original_filename)
+                documents = loader.load()
+                
+                # 更新进度
+                if progress_callback:
+                    progress_callback(file_size * 0.6, f"处理文件内容...")
+                    
+                # 使用原始文件名而不是存储的文件名
+                file_name = original_filename or os.path.basename(path)
+                for doc in documents:
+                    doc.metadata['file_name'] = file_name
+            except Exception as e:
+                print(f"加载文件时出错: {str(e)}")
+                raise
+        
+        # 分块
+        chunks = self.text_splitter.split_documents(documents)
+        
+        # 更新进度
+        if progress_callback:
+            if os.path.isdir(path):
+                progress_callback(processed_size, f"文档分块完成，共{len(chunks)}个文档片段")
+            else:
+                file_size = os.path.getsize(path)
+                progress_callback(file_size * 0.9, f"文档分块完成，共{len(chunks)}个文档片段")
+        
+        # 处理成统一格式
+        processed_docs = []
+        for i, chunk in enumerate(chunks):
+            processed_docs.append({
+                'id': f'doc_{i}',
+                'content': chunk.page_content,
+                'metadata': chunk.metadata
+            })
+            
         return processed_docs