4.8 KiB
4.8 KiB
文件处理功能说明
概述
BaseRAG 类现在支持自动处理多种格式的文件,并将它们转换为向量嵌入存储到知识库中。该功能包括:
- 支持的文件格式:txt、md、doc/docx
- 文件状态管理:使用 SQLite 数据库记录处理状态
- 文件存储管理:自动处理文件冲突和覆盖
- 简洁的API设计:易于使用和扩展
新增的类和功能
FileStatus 枚举
class FileStatus(Enum):
WAITING = "等待中"
PROCESSING = "处理中"
COMPLETED = "处理完毕"
ERROR = "处理失败"
FileManager 类
负责文件存储和状态管理:
- 文件哈希计算,避免重复处理
- SQLite 数据库记录文件状态
- 自动处理文件名冲突
BaseRAG 新增方法
process_file_to_vector_store(file_path, chunk_size=500, chunk_overlap=50)
主要的文件处理方法:
- 自动检测文件类型
- 保存文件到存储目录
- 切分文档并添加到向量库
- 记录处理状态
get_file_processing_status(file_hash=None)
获取文件处理状态:
- 传入 file_hash 获取特定文件状态
- 不传参数获取所有文件状态
list_files_by_status(status=None)
按状态筛选文件:
- 传入 FileStatus 枚举获取特定状态的文件
- 不传参数获取所有文件
使用示例
基本用法
from base_rag.core import BaseRAG, FileStatus
# 创建 RAG 实例
rag = SimpleRAG(
vector_store_name="my_knowledge_base",
storage_directory="./documents", # 文件存储目录
status_db_path="./file_status.db" # 状态数据库路径
)
# 处理文件
result = rag.process_file_to_vector_store("path/to/your/document.txt")
print(result)
# 查看处理状态
status = rag.get_file_processing_status()
print(status)
批量处理文件
import os
from pathlib import Path
# 处理目录中的所有文件
docs_dir = Path("./my_documents")
for file_path in docs_dir.glob("*"):
if file_path.suffix.lower() in ['.txt', '.md', '.doc', '.docx']:
print(f"处理文件: {file_path.name}")
result = rag.process_file_to_vector_store(str(file_path))
print(f"结果: {result['message']}")
# 查看处理结果统计
completed = rag.list_files_by_status(FileStatus.COMPLETED)
failed = rag.list_files_by_status(FileStatus.ERROR)
print(f"成功处理: {len(completed)} 个文件")
print(f"处理失败: {len(failed)} 个文件")
文件处理流程
- 文件保存:计算文件哈希,复制到存储目录,避免重复存储
- 状态检查:检查文件是否已经处理过,避免重复处理
- 状态更新:将状态更新为"等待中",然后"处理中"
- 文档加载:根据文件类型选择适当的加载器
- 文档切分:使用 RecursiveCharacterTextSplitter 切分文档
- 向量化:将文档片段转换为向量并存储
- 状态完成:更新状态为"处理完毕"或"处理失败"
数据库结构
SQLite 数据库表结构:
CREATE TABLE file_status (
id INTEGER PRIMARY KEY AUTOINCREMENT,
filename TEXT NOT NULL, -- 文件名
file_type TEXT NOT NULL, -- 文件类型 (.txt, .md, .doc 等)
file_hash TEXT UNIQUE NOT NULL, -- 文件哈希值
status TEXT NOT NULL, -- 处理状态
created_at TEXT NOT NULL, -- 创建时间
updated_at TEXT NOT NULL, -- 更新时间
error_message TEXT -- 错误信息(如果有)
);
配置选项
BaseRAG 初始化参数
BaseRAG(
vector_store_name="default", # 向量库名称
retriever_top_k=3, # 检索返回文档数
persist_directory="./chroma_db", # 向量库持久化目录
storage_directory="./documents", # 文件存储目录
status_db_path="./file_status.db", # 状态数据库路径
embedding_config=None, # 嵌入模型配置
rerank_config=None, # 重排模型配置
llm=None # 语言模型
)
文档切分参数
rag.process_file_to_vector_store(
file_path="document.txt",
chunk_size=500, # 切分块大小
chunk_overlap=50 # 切分重叠大小
)
安装依赖
确保安装了必要的依赖:
pip install unstructured python-docx
注意事项
- 文件冲突处理:相同内容的文件会被自动跳过,不会重复处理
- 错误处理:处理失败的文件状态会被记录,可以查看错误信息
- 性能考虑:大文件会被自动切分,建议根据实际情况调整 chunk_size
- 扩展性:可以通过重写
_load_document_by_type
方法支持更多文件格式
完整示例
参见 examples/file_processing_example.py
获取完整的使用示例。