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