feat: 文档处理"
This commit is contained in:
parent
5b101f26c6
commit
185344ee7a
|
@ -0,0 +1,160 @@
|
|||
# 文件处理功能说明
|
||||
|
||||
## 概述
|
||||
|
||||
BaseRAG 类现在支持自动处理多种格式的文件,并将它们转换为向量嵌入存储到知识库中。该功能包括:
|
||||
|
||||
1. **支持的文件格式**:txt、md、doc/docx
|
||||
2. **文件状态管理**:使用 SQLite 数据库记录处理状态
|
||||
3. **文件存储管理**:自动处理文件冲突和覆盖
|
||||
4. **简洁的API设计**:易于使用和扩展
|
||||
|
||||
## 新增的类和功能
|
||||
|
||||
### FileStatus 枚举
|
||||
```python
|
||||
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 枚举获取特定状态的文件
|
||||
- 不传参数获取所有文件
|
||||
|
||||
## 使用示例
|
||||
|
||||
### 基本用法
|
||||
|
||||
```python
|
||||
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)
|
||||
```
|
||||
|
||||
### 批量处理文件
|
||||
|
||||
```python
|
||||
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)} 个文件")
|
||||
```
|
||||
|
||||
## 文件处理流程
|
||||
|
||||
1. **文件保存**:计算文件哈希,复制到存储目录,避免重复存储
|
||||
2. **状态检查**:检查文件是否已经处理过,避免重复处理
|
||||
3. **状态更新**:将状态更新为"等待中",然后"处理中"
|
||||
4. **文档加载**:根据文件类型选择适当的加载器
|
||||
5. **文档切分**:使用 RecursiveCharacterTextSplitter 切分文档
|
||||
6. **向量化**:将文档片段转换为向量并存储
|
||||
7. **状态完成**:更新状态为"处理完毕"或"处理失败"
|
||||
|
||||
## 数据库结构
|
||||
|
||||
SQLite 数据库表结构:
|
||||
```sql
|
||||
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 初始化参数
|
||||
```python
|
||||
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 # 语言模型
|
||||
)
|
||||
```
|
||||
|
||||
### 文档切分参数
|
||||
```python
|
||||
rag.process_file_to_vector_store(
|
||||
file_path="document.txt",
|
||||
chunk_size=500, # 切分块大小
|
||||
chunk_overlap=50 # 切分重叠大小
|
||||
)
|
||||
```
|
||||
|
||||
## 安装依赖
|
||||
|
||||
确保安装了必要的依赖:
|
||||
```bash
|
||||
pip install unstructured python-docx
|
||||
```
|
||||
|
||||
## 注意事项
|
||||
|
||||
1. **文件冲突处理**:相同内容的文件会被自动跳过,不会重复处理
|
||||
2. **错误处理**:处理失败的文件状态会被记录,可以查看错误信息
|
||||
3. **性能考虑**:大文件会被自动切分,建议根据实际情况调整 chunk_size
|
||||
4. **扩展性**:可以通过重写 `_load_document_by_type` 方法支持更多文件格式
|
||||
|
||||
## 完整示例
|
||||
|
||||
参见 `examples/file_processing_example.py` 获取完整的使用示例。
|
|
@ -0,0 +1,148 @@
|
|||
# 快速开始指南
|
||||
|
||||
## 安装依赖
|
||||
|
||||
1. 激活虚拟环境:
|
||||
```bash
|
||||
source venv/bin/activate
|
||||
```
|
||||
|
||||
2. 安装依赖:
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## 基本使用
|
||||
|
||||
### 1. 创建 RAG 类实例
|
||||
|
||||
```python
|
||||
from base_rag.core import BaseRAG, FileStatus
|
||||
|
||||
class MyRAG(BaseRAG):
|
||||
def ingest(self, file_path: str, **kwargs):
|
||||
return self.process_file_to_vector_store(file_path, **kwargs)
|
||||
|
||||
def query(self, question: str) -> str:
|
||||
docs = self.similarity_search_with_rerank(question)
|
||||
if not docs:
|
||||
return "没有找到相关信息"
|
||||
return "\n".join([doc.page_content for doc in docs])
|
||||
|
||||
# 创建实例
|
||||
rag = MyRAG(
|
||||
vector_store_name="my_kb", # 知识库名称
|
||||
storage_directory="./documents", # 文件存储目录
|
||||
status_db_path="./file_status.db" # 状态数据库
|
||||
)
|
||||
```
|
||||
|
||||
### 2. 处理文件
|
||||
|
||||
```python
|
||||
# 处理单个文件
|
||||
result = rag.ingest("path/to/your/document.txt")
|
||||
print(f"处理结果: {result['message']}")
|
||||
|
||||
# 批量处理文件
|
||||
import os
|
||||
for filename in os.listdir("./documents"):
|
||||
if filename.endswith(('.txt', '.md', '.doc', '.docx')):
|
||||
result = rag.ingest(f"./documents/{filename}")
|
||||
print(f"{filename}: {result['message']}")
|
||||
```
|
||||
|
||||
### 3. 查询知识库
|
||||
|
||||
```python
|
||||
# 搜索相关文档
|
||||
answer = rag.query("你的问题")
|
||||
print(answer)
|
||||
```
|
||||
|
||||
### 4. 查看文件状态
|
||||
|
||||
```python
|
||||
# 查看所有文件状态
|
||||
all_files = rag.get_file_processing_status()
|
||||
for file_info in all_files:
|
||||
print(f"{file_info['filename']}: {file_info['status']}")
|
||||
|
||||
# 查看已完成的文件
|
||||
completed = rag.list_files_by_status(FileStatus.COMPLETED)
|
||||
print(f"已处理完成: {len(completed)} 个文件")
|
||||
|
||||
# 查看处理失败的文件
|
||||
failed = rag.list_files_by_status(FileStatus.ERROR)
|
||||
for file_info in failed:
|
||||
print(f"失败文件: {file_info['filename']}")
|
||||
print(f"错误信息: {file_info['error_message']}")
|
||||
```
|
||||
|
||||
## 支持的文件格式
|
||||
|
||||
- **.txt** - 纯文本文件
|
||||
- **.md** - Markdown 文件
|
||||
- **.doc/.docx** - Word 文档(需要安装 `unstructured` 和 `python-docx`)
|
||||
|
||||
## 主要特性
|
||||
|
||||
1. **自动去重**:相同内容的文件不会重复处理
|
||||
2. **状态跟踪**:实时跟踪文件处理状态
|
||||
3. **错误处理**:处理失败的文件会记录错误信息
|
||||
4. **简单API**:易于使用和扩展
|
||||
5. **持久化存储**:使用 SQLite 数据库记录状态
|
||||
|
||||
## 运行示例
|
||||
|
||||
```bash
|
||||
# 激活环境
|
||||
source venv/bin/activate
|
||||
|
||||
# 运行完整示例
|
||||
python examples/file_processing_example.py
|
||||
|
||||
# 运行简单测试
|
||||
python examples/simple_test.py
|
||||
```
|
||||
|
||||
## 配置选项
|
||||
|
||||
### 文档切分参数
|
||||
```python
|
||||
result = rag.ingest(
|
||||
"document.txt",
|
||||
chunk_size=500, # 切分块大小
|
||||
chunk_overlap=50 # 重叠大小
|
||||
)
|
||||
```
|
||||
|
||||
### 嵌入模型配置
|
||||
```python
|
||||
rag = MyRAG(
|
||||
embedding_config={
|
||||
"type": "local",
|
||||
"model_name": "BAAI/bge-small-zh-v1.5"
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
### 重排模型配置
|
||||
```python
|
||||
rag = MyRAG(
|
||||
rerank_config={
|
||||
"enabled": True,
|
||||
"type": "local",
|
||||
"model": "BAAI/bge-reranker-base",
|
||||
"top_k": 3
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
## 数据存储
|
||||
|
||||
- **文件存储**:`./documents/` 目录(可配置)
|
||||
- **向量数据库**:`./chroma_db/` 目录
|
||||
- **状态数据库**:`./file_status.db` 文件
|
||||
|
||||
文件名格式:`原文件名_哈希值前8位.扩展名`
|
|
@ -0,0 +1,6 @@
|
|||
|
||||
这是一个测试文档。
|
||||
它包含了关于人工智能的信息。
|
||||
人工智能是计算机科学的一个分支,致力于创建智能机器。
|
||||
机器学习是人工智能的一个重要组成部分。
|
||||
深度学习是机器学习的一个子领域。
|
|
@ -0,0 +1,13 @@
|
|||
|
||||
# RAG系统介绍
|
||||
|
||||
## 什么是RAG?
|
||||
RAG(Retrieval-Augmented Generation)是一种结合了检索和生成的AI技术。
|
||||
|
||||
## RAG的优势
|
||||
- 能够利用外部知识库
|
||||
- 提高回答的准确性
|
||||
- 支持实时更新知识
|
||||
|
||||
## 应用场景
|
||||
RAG系统广泛应用于问答系统、知识管理等领域。
|
Binary file not shown.
|
@ -0,0 +1,126 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
文件处理示例
|
||||
演示如何使用BaseRAG的文件处理功能
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# 添加源码路径
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'src'))
|
||||
|
||||
from base_rag.core import BaseRAG, FileStatus
|
||||
|
||||
|
||||
class SimpleRAG(BaseRAG):
|
||||
"""简单的RAG实现示例"""
|
||||
|
||||
def ingest(self, file_path: str, **kwargs):
|
||||
"""实现文档导入逻辑"""
|
||||
return self.process_file_to_vector_store(file_path, **kwargs)
|
||||
|
||||
def query(self, question: str) -> str:
|
||||
"""实现简单的查询逻辑"""
|
||||
# 使用相似性搜索
|
||||
docs = self.similarity_search_with_rerank(question)
|
||||
|
||||
if not docs:
|
||||
return "抱歉,没有找到相关信息。"
|
||||
|
||||
# 简单的回答拼接(实际应用中应该使用LLM)
|
||||
context = "\n".join([doc.page_content for doc in docs])
|
||||
return f"基于以下信息回答:\n{context}"
|
||||
|
||||
|
||||
def main():
|
||||
# 创建RAG实例
|
||||
rag = SimpleRAG(
|
||||
vector_store_name="file_demo",
|
||||
retriever_top_k=3,
|
||||
storage_directory="./demo_documents",
|
||||
status_db_path="./demo_file_status.db"
|
||||
)
|
||||
|
||||
# 创建测试文件
|
||||
test_dir = Path("./test_files")
|
||||
test_dir.mkdir(exist_ok=True)
|
||||
|
||||
# 创建测试文本文件
|
||||
txt_file = test_dir / "test_document.txt"
|
||||
txt_file.write_text("""
|
||||
这是一个测试文档。
|
||||
它包含了关于人工智能的信息。
|
||||
人工智能是计算机科学的一个分支,致力于创建智能机器。
|
||||
机器学习是人工智能的一个重要组成部分。
|
||||
深度学习是机器学习的一个子领域。
|
||||
""", encoding="utf-8")
|
||||
|
||||
# 创建测试Markdown文件
|
||||
md_file = test_dir / "test_markdown.md"
|
||||
md_file.write_text("""
|
||||
# RAG系统介绍
|
||||
|
||||
## 什么是RAG?
|
||||
RAG(Retrieval-Augmented Generation)是一种结合了检索和生成的AI技术。
|
||||
|
||||
## RAG的优势
|
||||
- 能够利用外部知识库
|
||||
- 提高回答的准确性
|
||||
- 支持实时更新知识
|
||||
|
||||
## 应用场景
|
||||
RAG系统广泛应用于问答系统、知识管理等领域。
|
||||
""", encoding="utf-8")
|
||||
|
||||
print("=== 文件处理示例 ===\n")
|
||||
|
||||
# 1. 处理文本文件
|
||||
print("1. 处理文本文件...")
|
||||
result1 = rag.ingest(str(txt_file))
|
||||
print(f"处理结果: {result1}\n")
|
||||
|
||||
# 2. 处理Markdown文件
|
||||
print("2. 处理Markdown文件...")
|
||||
result2 = rag.ingest(str(md_file))
|
||||
print(f"处理结果: {result2}\n")
|
||||
|
||||
# 3. 再次处理同一个文件(应该跳过)
|
||||
print("3. 再次处理文本文件(测试重复处理)...")
|
||||
result3 = rag.ingest(str(txt_file))
|
||||
print(f"处理结果: {result3}\n")
|
||||
|
||||
# 4. 查看所有文件状态
|
||||
print("4. 查看所有文件处理状态...")
|
||||
all_files = rag.get_file_processing_status()
|
||||
for file_info in all_files:
|
||||
print(f"文件: {file_info['filename']}")
|
||||
print(f"类型: {file_info['file_type']}")
|
||||
print(f"状态: {file_info['status']}")
|
||||
print(f"处理时间: {file_info['updated_at']}")
|
||||
print("---")
|
||||
|
||||
# 5. 查看已完成的文件
|
||||
print("\n5. 查看已完成处理的文件...")
|
||||
completed_files = rag.list_files_by_status(FileStatus.COMPLETED)
|
||||
print(f"已完成处理的文件数量: {len(completed_files)}")
|
||||
|
||||
# 6. 测试搜索功能
|
||||
print("\n6. 测试搜索功能...")
|
||||
questions = [
|
||||
"什么是人工智能?",
|
||||
"RAG有什么优势?",
|
||||
"机器学习是什么?"
|
||||
]
|
||||
|
||||
for question in questions:
|
||||
print(f"\n问题: {question}")
|
||||
answer = rag.query(question)
|
||||
print(f"回答: {answer[:200]}...") # 只显示前200个字符
|
||||
|
||||
print("\n=== 示例完成 ===")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
|
@ -0,0 +1,88 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
简单的文件处理测试
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# 添加源码路径
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'src'))
|
||||
|
||||
from base_rag.core import BaseRAG, FileStatus
|
||||
|
||||
|
||||
class SimpleRAG(BaseRAG):
|
||||
"""简单的RAG实现示例"""
|
||||
|
||||
def ingest(self, file_path: str, **kwargs):
|
||||
"""实现文档导入逻辑"""
|
||||
return self.process_file_to_vector_store(file_path, **kwargs)
|
||||
|
||||
def query(self, question: str) -> str:
|
||||
"""实现简单的查询逻辑"""
|
||||
docs = self.similarity_search_with_rerank(question)
|
||||
|
||||
if not docs:
|
||||
return "抱歉,没有找到相关信息。"
|
||||
|
||||
# 简单的回答拼接
|
||||
context = "\n".join([doc.page_content for doc in docs])
|
||||
return f"基于以下信息回答:\n{context}"
|
||||
|
||||
|
||||
def test_file_processing():
|
||||
print("=== 文件处理功能测试 ===\n")
|
||||
|
||||
# 创建RAG实例
|
||||
rag = SimpleRAG(
|
||||
vector_store_name="test_kb",
|
||||
retriever_top_k=2,
|
||||
storage_directory="./test_docs",
|
||||
status_db_path="./test_status.db"
|
||||
)
|
||||
|
||||
# 创建测试文件
|
||||
test_dir = Path("./test_files")
|
||||
test_dir.mkdir(exist_ok=True)
|
||||
|
||||
# 创建一个知识文件
|
||||
knowledge_file = test_dir / "knowledge.txt"
|
||||
knowledge_file.write_text("""
|
||||
Python是一种高级编程语言。
|
||||
它具有简洁的语法和强大的功能。
|
||||
Python广泛应用于Web开发、数据科学、人工智能等领域。
|
||||
机器学习库如scikit-learn、TensorFlow和PyTorch都支持Python。
|
||||
Flask和Django是流行的Python Web框架。
|
||||
""", encoding="utf-8")
|
||||
|
||||
print("1. 处理知识文件...")
|
||||
result = rag.ingest(str(knowledge_file))
|
||||
print(f"处理结果: {result['message']}")
|
||||
print(f"文档片段数: {result.get('chunks_count', 0)}")
|
||||
print()
|
||||
|
||||
print("2. 查询测试...")
|
||||
questions = [
|
||||
"Python是什么?",
|
||||
"Python有哪些应用领域?",
|
||||
"有哪些Python Web框架?"
|
||||
]
|
||||
|
||||
for question in questions:
|
||||
print(f"问题: {question}")
|
||||
answer = rag.query(question)
|
||||
print(f"回答: {answer[:100]}...")
|
||||
print()
|
||||
|
||||
print("3. 查看文件状态...")
|
||||
files = rag.get_file_processing_status()
|
||||
for file_info in files:
|
||||
print(f"文件: {file_info['filename']} | 状态: {file_info['status']}")
|
||||
|
||||
print("\n=== 测试完成 ===")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_file_processing()
|
|
@ -6,6 +6,10 @@ chromadb>=0.4.0
|
|||
sentence-transformers>=2.2.0
|
||||
numpy>=1.21.0
|
||||
|
||||
# 文档处理依赖
|
||||
unstructured>=0.10.0
|
||||
python-docx>=0.8.11
|
||||
|
||||
# 可选依赖
|
||||
# langchain-openai>=0.2.0 # 用于本地API接口支持
|
||||
# FlagEmbedding>=1.2.0 # 用于BGE重排
|
|
@ -2,6 +2,13 @@ from abc import ABC, abstractmethod
|
|||
from typing import List, Optional, Dict, ClassVar, Union, Tuple, Any
|
||||
import threading
|
||||
import numpy as np
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import hashlib
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from enum import Enum
|
||||
|
||||
from langchain_huggingface import HuggingFaceEmbeddings
|
||||
from langchain.embeddings.base import Embeddings
|
||||
|
@ -13,6 +20,151 @@ from langchain.llms.base import BaseLLM
|
|||
from langchain.schema import Document
|
||||
|
||||
|
||||
class FileStatus(Enum):
|
||||
"""文件处理状态枚举"""
|
||||
WAITING = "等待中"
|
||||
PROCESSING = "处理中"
|
||||
COMPLETED = "处理完毕"
|
||||
ERROR = "处理失败"
|
||||
|
||||
|
||||
class FileManager:
|
||||
"""文件管理器,负责文件存储、状态记录等"""
|
||||
|
||||
def __init__(self, storage_dir: str = "./documents", db_path: str = "./file_status.db"):
|
||||
self.storage_dir = Path(storage_dir)
|
||||
self.db_path = db_path
|
||||
self.storage_dir.mkdir(exist_ok=True)
|
||||
self._init_database()
|
||||
|
||||
def _init_database(self):
|
||||
"""初始化状态记录数据库"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS file_status (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
filename TEXT NOT NULL,
|
||||
file_type TEXT NOT NULL,
|
||||
file_hash TEXT UNIQUE NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
error_message TEXT
|
||||
)
|
||||
""")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
def _calculate_file_hash(self, file_path: str) -> str:
|
||||
"""计算文件哈希值"""
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(file_path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_md5.update(chunk)
|
||||
return hash_md5.hexdigest()
|
||||
|
||||
def save_file(self, source_path: str) -> Tuple[str, str]:
|
||||
"""
|
||||
保存文件到存储目录
|
||||
返回: (存储路径, 文件哈希)
|
||||
"""
|
||||
source_path = Path(source_path)
|
||||
if not source_path.exists():
|
||||
raise FileNotFoundError(f"源文件不存在: {source_path}")
|
||||
|
||||
# 计算文件哈希
|
||||
file_hash = self._calculate_file_hash(str(source_path))
|
||||
|
||||
# 生成存储文件名(使用哈希前8位避免冲突)
|
||||
file_extension = source_path.suffix
|
||||
stored_filename = f"{source_path.stem}_{file_hash[:8]}{file_extension}"
|
||||
stored_path = self.storage_dir / stored_filename
|
||||
|
||||
# 如果文件已存在且哈希相同,直接返回
|
||||
if stored_path.exists():
|
||||
existing_hash = self._calculate_file_hash(str(stored_path))
|
||||
if existing_hash == file_hash:
|
||||
print(f"文件已存在,跳过复制: {stored_filename}")
|
||||
return str(stored_path), file_hash
|
||||
|
||||
# 复制文件
|
||||
shutil.copy2(source_path, stored_path)
|
||||
print(f"文件已保存到: {stored_path}")
|
||||
|
||||
return str(stored_path), file_hash
|
||||
|
||||
def update_file_status(self, file_hash: str, filename: str, file_type: str,
|
||||
status: FileStatus, error_message: str = None):
|
||||
"""更新文件处理状态"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
now = datetime.now().isoformat()
|
||||
|
||||
# 尝试更新现有记录
|
||||
cursor.execute("""
|
||||
UPDATE file_status
|
||||
SET status = ?, updated_at = ?, error_message = ?
|
||||
WHERE file_hash = ?
|
||||
""", (status.value, now, error_message, file_hash))
|
||||
|
||||
# 如果没有更新任何记录,插入新记录
|
||||
if cursor.rowcount == 0:
|
||||
cursor.execute("""
|
||||
INSERT INTO file_status (filename, file_type, file_hash, status, created_at, updated_at, error_message)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
""", (filename, file_type, file_hash, status.value, now, now, error_message))
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
def get_file_status(self, file_hash: str) -> Optional[Dict]:
|
||||
"""获取文件状态"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM file_status WHERE file_hash = ?", (file_hash,))
|
||||
row = cursor.fetchone()
|
||||
conn.close()
|
||||
|
||||
if row:
|
||||
return {
|
||||
'id': row[0],
|
||||
'filename': row[1],
|
||||
'file_type': row[2],
|
||||
'file_hash': row[3],
|
||||
'status': row[4],
|
||||
'created_at': row[5],
|
||||
'updated_at': row[6],
|
||||
'error_message': row[7]
|
||||
}
|
||||
return None
|
||||
|
||||
def list_files_by_status(self, status: FileStatus = None) -> List[Dict]:
|
||||
"""列出指定状态的文件"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
if status:
|
||||
cursor.execute("SELECT * FROM file_status WHERE status = ? ORDER BY created_at DESC", (status.value,))
|
||||
else:
|
||||
cursor.execute("SELECT * FROM file_status ORDER BY created_at DESC")
|
||||
|
||||
rows = cursor.fetchall()
|
||||
conn.close()
|
||||
|
||||
return [{
|
||||
'id': row[0],
|
||||
'filename': row[1],
|
||||
'file_type': row[2],
|
||||
'file_hash': row[3],
|
||||
'status': row[4],
|
||||
'created_at': row[5],
|
||||
'updated_at': row[6],
|
||||
'error_message': row[7]
|
||||
} for row in rows]
|
||||
|
||||
|
||||
class ModelManager:
|
||||
"""统一的模型管理类,用于创建和缓存embedding和rerank模型"""
|
||||
|
||||
|
@ -182,6 +334,8 @@ class BaseRAG(ABC):
|
|||
llm: Optional[BaseLLM] = None,
|
||||
embedding_config: Optional[Dict] = None,
|
||||
rerank_config: Optional[Dict] = None,
|
||||
storage_directory: str = "./documents",
|
||||
status_db_path: str = "./file_status.db",
|
||||
):
|
||||
"""
|
||||
初始化基础RAG类。
|
||||
|
@ -191,6 +345,8 @@ class BaseRAG(ABC):
|
|||
:param llm: 可选的对话模型
|
||||
:param persist_directory: Chroma持久化目录
|
||||
:param rerank_config: 重排配置
|
||||
:param storage_directory: 文件存储目录
|
||||
:param status_db_path: 文件状态数据库路径
|
||||
|
||||
embedding_config 示例:
|
||||
本地模型名称: {"type": "local", "model_name": "BAAI/bge-small-zh-v1.5"}
|
||||
|
@ -212,6 +368,9 @@ class BaseRAG(ABC):
|
|||
self.persist_directory = persist_directory
|
||||
self.rerank_config = rerank_config or {"enabled": False}
|
||||
|
||||
# 初始化文件管理器
|
||||
self.file_manager = FileManager(storage_directory, status_db_path)
|
||||
|
||||
# 使用统一的模型管理器创建嵌入模型
|
||||
self.embedding_model = ModelManager.get_or_create_model(
|
||||
self.embedding_config, "embedding", ModelManager.create_embedding_model
|
||||
|
@ -343,13 +502,168 @@ class BaseRAG(ABC):
|
|||
splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=20)
|
||||
return splitter.split_documents(documents)
|
||||
|
||||
def _load_document_by_type(self, file_path: str) -> List[Document]:
|
||||
"""
|
||||
根据文件类型加载文档
|
||||
"""
|
||||
file_path = Path(file_path)
|
||||
file_extension = file_path.suffix.lower()
|
||||
|
||||
try:
|
||||
if file_extension in ['.txt', '.md']:
|
||||
# 文本和Markdown文件
|
||||
loader = TextLoader(str(file_path), encoding="utf-8")
|
||||
return loader.load()
|
||||
|
||||
elif file_extension in ['.doc', '.docx']:
|
||||
# Word文档
|
||||
try:
|
||||
from langchain_community.document_loaders import UnstructuredWordDocumentLoader
|
||||
loader = UnstructuredWordDocumentLoader(str(file_path))
|
||||
return loader.load()
|
||||
except ImportError:
|
||||
print("警告: 需要安装 unstructured 和 python-docx 来处理Word文档")
|
||||
print("请运行: pip install unstructured python-docx")
|
||||
raise
|
||||
|
||||
else:
|
||||
raise ValueError(f"不支持的文件类型: {file_extension}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"加载文件失败 {file_path}: {e}")
|
||||
raise
|
||||
|
||||
def process_file_to_vector_store(self, file_path: str, chunk_size: int = 500, chunk_overlap: int = 50) -> Dict:
|
||||
"""
|
||||
处理文件并添加到向量库
|
||||
|
||||
:param file_path: 文件路径
|
||||
:param chunk_size: 文档切分大小
|
||||
:param chunk_overlap: 文档切分重叠
|
||||
:return: 处理结果字典
|
||||
"""
|
||||
file_path = Path(file_path)
|
||||
if not file_path.exists():
|
||||
raise FileNotFoundError(f"文件不存在: {file_path}")
|
||||
|
||||
filename = file_path.name
|
||||
file_type = file_path.suffix.lower()
|
||||
|
||||
try:
|
||||
# 1. 保存文件并获取哈希
|
||||
stored_path, file_hash = self.file_manager.save_file(str(file_path))
|
||||
|
||||
# 2. 检查文件是否已经处理过
|
||||
existing_status = self.file_manager.get_file_status(file_hash)
|
||||
if existing_status and existing_status['status'] == FileStatus.COMPLETED.value:
|
||||
print(f"文件 {filename} 已经处理完毕,跳过处理")
|
||||
return {
|
||||
'success': True,
|
||||
'message': '文件已存在,跳过处理',
|
||||
'file_hash': file_hash,
|
||||
'filename': filename,
|
||||
'status': FileStatus.COMPLETED.value
|
||||
}
|
||||
|
||||
# 3. 更新状态为等待中
|
||||
self.file_manager.update_file_status(
|
||||
file_hash, filename, file_type, FileStatus.WAITING
|
||||
)
|
||||
|
||||
# 4. 更新状态为处理中
|
||||
self.file_manager.update_file_status(
|
||||
file_hash, filename, file_type, FileStatus.PROCESSING
|
||||
)
|
||||
|
||||
# 5. 加载文档
|
||||
print(f"开始处理文件: {filename}")
|
||||
documents = self._load_document_by_type(stored_path)
|
||||
|
||||
if not documents:
|
||||
raise ValueError("未能从文件中提取到任何内容")
|
||||
|
||||
# 6. 切分文档
|
||||
splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap
|
||||
)
|
||||
split_docs = splitter.split_documents(documents)
|
||||
|
||||
# 7. 为每个切分的文档添加元数据
|
||||
for doc in split_docs:
|
||||
doc.metadata.update({
|
||||
'source_file': filename,
|
||||
'file_hash': file_hash,
|
||||
'file_type': file_type,
|
||||
'processed_at': datetime.now().isoformat()
|
||||
})
|
||||
|
||||
# 8. 添加到向量库
|
||||
print(f"将 {len(split_docs)} 个文档片段添加到向量库...")
|
||||
self.add_documents_to_vector_store(split_docs)
|
||||
|
||||
# 9. 更新状态为完成
|
||||
self.file_manager.update_file_status(
|
||||
file_hash, filename, file_type, FileStatus.COMPLETED
|
||||
)
|
||||
|
||||
print(f"文件处理完成: {filename} ({len(split_docs)} 个片段)")
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'message': '文件处理完成',
|
||||
'file_hash': file_hash,
|
||||
'filename': filename,
|
||||
'chunks_count': len(split_docs),
|
||||
'status': FileStatus.COMPLETED.value
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
print(f"文件处理失败 {filename}: {error_message}")
|
||||
|
||||
# 更新状态为错误
|
||||
if 'file_hash' in locals():
|
||||
self.file_manager.update_file_status(
|
||||
file_hash, filename, file_type, FileStatus.ERROR, error_message
|
||||
)
|
||||
|
||||
return {
|
||||
'success': False,
|
||||
'message': f'文件处理失败: {error_message}',
|
||||
'filename': filename,
|
||||
'status': FileStatus.ERROR.value,
|
||||
'error': error_message
|
||||
}
|
||||
|
||||
def get_file_processing_status(self, file_hash: str = None) -> Union[Dict, List[Dict]]:
|
||||
"""
|
||||
获取文件处理状态
|
||||
|
||||
:param file_hash: 文件哈希,如果为None则返回所有文件状态
|
||||
:return: 文件状态信息
|
||||
"""
|
||||
if file_hash:
|
||||
return self.file_manager.get_file_status(file_hash)
|
||||
else:
|
||||
return self.file_manager.list_files_by_status()
|
||||
|
||||
def list_files_by_status(self, status: FileStatus = None) -> List[Dict]:
|
||||
"""
|
||||
按状态列出文件
|
||||
|
||||
:param status: 文件状态,如果为None则返回所有状态的文件
|
||||
:return: 文件列表
|
||||
"""
|
||||
return self.file_manager.list_files_by_status(status)
|
||||
|
||||
def add_documents_to_vector_store(self, documents: List[Document]):
|
||||
"""
|
||||
将文档添加到 Chroma 向量库。
|
||||
"""
|
||||
if documents:
|
||||
self.vector_store.add_documents(documents)
|
||||
self.vector_store.persist() # 持久化数据
|
||||
# 新版本的 Chroma 会自动持久化数据
|
||||
|
||||
def build_retriever(self):
|
||||
"""
|
||||
|
|
|
@ -0,0 +1,6 @@
|
|||
|
||||
Python是一种高级编程语言。
|
||||
它具有简洁的语法和强大的功能。
|
||||
Python广泛应用于Web开发、数据科学、人工智能等领域。
|
||||
机器学习库如scikit-learn、TensorFlow和PyTorch都支持Python。
|
||||
Flask和Django是流行的Python Web框架。
|
|
@ -0,0 +1,6 @@
|
|||
|
||||
Python是一种高级编程语言。
|
||||
它具有简洁的语法和强大的功能。
|
||||
Python广泛应用于Web开发、数据科学、人工智能等领域。
|
||||
机器学习库如scikit-learn、TensorFlow和PyTorch都支持Python。
|
||||
Flask和Django是流行的Python Web框架。
|
|
@ -0,0 +1,6 @@
|
|||
|
||||
这是一个测试文档。
|
||||
它包含了关于人工智能的信息。
|
||||
人工智能是计算机科学的一个分支,致力于创建智能机器。
|
||||
机器学习是人工智能的一个重要组成部分。
|
||||
深度学习是机器学习的一个子领域。
|
|
@ -0,0 +1,13 @@
|
|||
|
||||
# RAG系统介绍
|
||||
|
||||
## 什么是RAG?
|
||||
RAG(Retrieval-Augmented Generation)是一种结合了检索和生成的AI技术。
|
||||
|
||||
## RAG的优势
|
||||
- 能够利用外部知识库
|
||||
- 提高回答的准确性
|
||||
- 支持实时更新知识
|
||||
|
||||
## 应用场景
|
||||
RAG系统广泛应用于问答系统、知识管理等领域。
|
Binary file not shown.
Loading…
Reference in New Issue