feat: 文档处理"

This commit is contained in:
李如威 2025-08-07 16:14:21 +08:00
parent 5b101f26c6
commit 185344ee7a
14 changed files with 891 additions and 1 deletions

160
FILE_PROCESSING_GUIDE.md Normal file
View File

@ -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` 获取完整的使用示例。

148
QUICK_START.md Normal file
View File

@ -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位.扩展名`

View File

@ -0,0 +1,6 @@
这是一个测试文档。
它包含了关于人工智能的信息。
人工智能是计算机科学的一个分支,致力于创建智能机器。
机器学习是人工智能的一个重要组成部分。
深度学习是机器学习的一个子领域。

View File

@ -0,0 +1,13 @@
# RAG系统介绍
## 什么是RAG
RAGRetrieval-Augmented Generation是一种结合了检索和生成的AI技术。
## RAG的优势
- 能够利用外部知识库
- 提高回答的准确性
- 支持实时更新知识
## 应用场景
RAG系统广泛应用于问答系统、知识管理等领域。

BIN
demo_file_status.db Normal file

Binary file not shown.

View File

@ -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
RAGRetrieval-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()

88
examples/simple_test.py Normal file
View File

@ -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-learnTensorFlow和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()

View File

@ -6,6 +6,10 @@ chromadb>=0.4.0
sentence-transformers>=2.2.0 sentence-transformers>=2.2.0
numpy>=1.21.0 numpy>=1.21.0
# 文档处理依赖
unstructured>=0.10.0
python-docx>=0.8.11
# 可选依赖 # 可选依赖
# langchain-openai>=0.2.0 # 用于本地API接口支持 # langchain-openai>=0.2.0 # 用于本地API接口支持
# FlagEmbedding>=1.2.0 # 用于BGE重排 # FlagEmbedding>=1.2.0 # 用于BGE重排

View File

@ -2,6 +2,13 @@ from abc import ABC, abstractmethod
from typing import List, Optional, Dict, ClassVar, Union, Tuple, Any from typing import List, Optional, Dict, ClassVar, Union, Tuple, Any
import threading import threading
import numpy as np 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_huggingface import HuggingFaceEmbeddings
from langchain.embeddings.base import Embeddings from langchain.embeddings.base import Embeddings
@ -13,6 +20,151 @@ from langchain.llms.base import BaseLLM
from langchain.schema import Document 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: class ModelManager:
"""统一的模型管理类用于创建和缓存embedding和rerank模型""" """统一的模型管理类用于创建和缓存embedding和rerank模型"""
@ -182,6 +334,8 @@ class BaseRAG(ABC):
llm: Optional[BaseLLM] = None, llm: Optional[BaseLLM] = None,
embedding_config: Optional[Dict] = None, embedding_config: Optional[Dict] = None,
rerank_config: Optional[Dict] = None, rerank_config: Optional[Dict] = None,
storage_directory: str = "./documents",
status_db_path: str = "./file_status.db",
): ):
""" """
初始化基础RAG类 初始化基础RAG类
@ -191,6 +345,8 @@ class BaseRAG(ABC):
:param llm: 可选的对话模型 :param llm: 可选的对话模型
:param persist_directory: Chroma持久化目录 :param persist_directory: Chroma持久化目录
:param rerank_config: 重排配置 :param rerank_config: 重排配置
:param storage_directory: 文件存储目录
:param status_db_path: 文件状态数据库路径
embedding_config 示例: embedding_config 示例:
本地模型名称: {"type": "local", "model_name": "BAAI/bge-small-zh-v1.5"} 本地模型名称: {"type": "local", "model_name": "BAAI/bge-small-zh-v1.5"}
@ -212,6 +368,9 @@ class BaseRAG(ABC):
self.persist_directory = persist_directory self.persist_directory = persist_directory
self.rerank_config = rerank_config or {"enabled": False} 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_model = ModelManager.get_or_create_model(
self.embedding_config, "embedding", ModelManager.create_embedding_model self.embedding_config, "embedding", ModelManager.create_embedding_model
@ -343,13 +502,168 @@ class BaseRAG(ABC):
splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=20) splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=20)
return splitter.split_documents(documents) 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]): def add_documents_to_vector_store(self, documents: List[Document]):
""" """
将文档添加到 Chroma 向量库 将文档添加到 Chroma 向量库
""" """
if documents: if documents:
self.vector_store.add_documents(documents) self.vector_store.add_documents(documents)
self.vector_store.persist() # 持久化数据 # 新版本的 Chroma 会自动持久化数据
def build_retriever(self): def build_retriever(self):
""" """

View File

@ -0,0 +1,6 @@
Python是一种高级编程语言。
它具有简洁的语法和强大的功能。
Python广泛应用于Web开发、数据科学、人工智能等领域。
机器学习库如scikit-learn、TensorFlow和PyTorch都支持Python。
Flask和Django是流行的Python Web框架。

6
test_files/knowledge.txt Normal file
View File

@ -0,0 +1,6 @@
Python是一种高级编程语言。
它具有简洁的语法和强大的功能。
Python广泛应用于Web开发、数据科学、人工智能等领域。
机器学习库如scikit-learn、TensorFlow和PyTorch都支持Python。
Flask和Django是流行的Python Web框架。

View File

@ -0,0 +1,6 @@
这是一个测试文档。
它包含了关于人工智能的信息。
人工智能是计算机科学的一个分支,致力于创建智能机器。
机器学习是人工智能的一个重要组成部分。
深度学习是机器学习的一个子领域。

View File

@ -0,0 +1,13 @@
# RAG系统介绍
## 什么是RAG
RAGRetrieval-Augmented Generation是一种结合了检索和生成的AI技术。
## RAG的优势
- 能够利用外部知识库
- 提高回答的准确性
- 支持实时更新知识
## 应用场景
RAG系统广泛应用于问答系统、知识管理等领域。

BIN
test_status.db Normal file

Binary file not shown.