import sys import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from base_rag import BaseRAG class SimpleRAG(BaseRAG): def ingest(self, documents): for doc in documents: self.vector_store.add_texts([doc]) def query(self, question, k=3): docs = self.vector_store.similarity_search(question, k=k) return docs def main(): config = { "model_name": "sentence-transformers/all-MiniLM-L6-v2", "embedding_type": "local", } rag = SimpleRAG(embedding_config=config) print("RAG初始化完成!") # 添加一些文档 documents = [ "苹果是一种水果,味道甜美,营养丰富。", "苹果公司是一家科技公司,生产iPhone和Mac等产品。", "Python是一种编程语言,简单易学,功能强大。", ] print("正在添加文档...") rag.ingest(documents) print("文档添加完成!") # 测试查询 print("\n正在查询: '什么是苹果?'") result = rag.query("什么是苹果?") print(f"查询结果: {result}") if __name__ == "__main__": main()