diff --git a/FILE_PROCESSING_GUIDE.md b/FILE_PROCESSING_GUIDE.md new file mode 100644 index 0000000..d77bf1b --- /dev/null +++ b/FILE_PROCESSING_GUIDE.md @@ -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` 获取完整的使用示例。 diff --git a/QUICK_START.md b/QUICK_START.md new file mode 100644 index 0000000..fa09186 --- /dev/null +++ b/QUICK_START.md @@ -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位.扩展名` diff --git a/demo_documents/test_document_35daadf0.txt b/demo_documents/test_document_35daadf0.txt new file mode 100644 index 0000000..fbab07a --- /dev/null +++ b/demo_documents/test_document_35daadf0.txt @@ -0,0 +1,6 @@ + +这是一个测试文档。 +它包含了关于人工智能的信息。 +人工智能是计算机科学的一个分支,致力于创建智能机器。 +机器学习是人工智能的一个重要组成部分。 +深度学习是机器学习的一个子领域。 diff --git a/demo_documents/test_markdown_787e71fb.md b/demo_documents/test_markdown_787e71fb.md new file mode 100644 index 0000000..9bbb550 --- /dev/null +++ b/demo_documents/test_markdown_787e71fb.md @@ -0,0 +1,13 @@ + +# RAG系统介绍 + +## 什么是RAG? +RAG(Retrieval-Augmented Generation)是一种结合了检索和生成的AI技术。 + +## RAG的优势 +- 能够利用外部知识库 +- 提高回答的准确性 +- 支持实时更新知识 + +## 应用场景 +RAG系统广泛应用于问答系统、知识管理等领域。 diff --git a/demo_file_status.db b/demo_file_status.db new file mode 100644 index 0000000..0476d0e Binary files /dev/null and b/demo_file_status.db differ diff --git a/examples/file_processing_example.py b/examples/file_processing_example.py new file mode 100644 index 0000000..560396f --- /dev/null +++ b/examples/file_processing_example.py @@ -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() diff --git a/examples/simple_test.py b/examples/simple_test.py new file mode 100644 index 0000000..6d366dc --- /dev/null +++ b/examples/simple_test.py @@ -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() diff --git a/requirements.txt b/requirements.txt index 563529d..2277f59 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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重排 \ No newline at end of file diff --git a/src/base_rag/core.py b/src/base_rag/core.py index f839ff9..249da5a 100644 --- a/src/base_rag/core.py +++ b/src/base_rag/core.py @@ -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): """ diff --git a/test_docs/knowledge_e15fef20.txt b/test_docs/knowledge_e15fef20.txt new file mode 100644 index 0000000..fbb53a2 --- /dev/null +++ b/test_docs/knowledge_e15fef20.txt @@ -0,0 +1,6 @@ + +Python是一种高级编程语言。 +它具有简洁的语法和强大的功能。 +Python广泛应用于Web开发、数据科学、人工智能等领域。 +机器学习库如scikit-learn、TensorFlow和PyTorch都支持Python。 +Flask和Django是流行的Python Web框架。 diff --git a/test_files/knowledge.txt b/test_files/knowledge.txt new file mode 100644 index 0000000..fbb53a2 --- /dev/null +++ b/test_files/knowledge.txt @@ -0,0 +1,6 @@ + +Python是一种高级编程语言。 +它具有简洁的语法和强大的功能。 +Python广泛应用于Web开发、数据科学、人工智能等领域。 +机器学习库如scikit-learn、TensorFlow和PyTorch都支持Python。 +Flask和Django是流行的Python Web框架。 diff --git a/test_files/test_document.txt b/test_files/test_document.txt new file mode 100644 index 0000000..fbab07a --- /dev/null +++ b/test_files/test_document.txt @@ -0,0 +1,6 @@ + +这是一个测试文档。 +它包含了关于人工智能的信息。 +人工智能是计算机科学的一个分支,致力于创建智能机器。 +机器学习是人工智能的一个重要组成部分。 +深度学习是机器学习的一个子领域。 diff --git a/test_files/test_markdown.md b/test_files/test_markdown.md new file mode 100644 index 0000000..9bbb550 --- /dev/null +++ b/test_files/test_markdown.md @@ -0,0 +1,13 @@ + +# RAG系统介绍 + +## 什么是RAG? +RAG(Retrieval-Augmented Generation)是一种结合了检索和生成的AI技术。 + +## RAG的优势 +- 能够利用外部知识库 +- 提高回答的准确性 +- 支持实时更新知识 + +## 应用场景 +RAG系统广泛应用于问答系统、知识管理等领域。 diff --git a/test_status.db b/test_status.db new file mode 100644 index 0000000..660fc60 Binary files /dev/null and b/test_status.db differ