127 lines
3.7 KiB
Python
127 lines
3.7 KiB
Python
#!/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()
|