47 lines
1.1 KiB
Python
47 lines
1.1 KiB
Python
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()
|