base_rag/examples/file_processing_example.py

127 lines
3.7 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/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()