39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
from rag_system import RAGSystem
|
||
|
||
|
||
def main():
|
||
# 初始化RAG系统
|
||
rag_system = RAGSystem()
|
||
|
||
# 训练文档
|
||
# documents = [
|
||
# {
|
||
# 'content': '''Reisasol(零咲)是一Reisa的女儿，Rating高达15684。外貌设定为白色头发、红色瞳孔，性格是带有性压抑倾向的瑟琴本子女主，内心缺爱、渴望被关注与亲密接触，外表清纯但内在有阴暗与矛盾感。''',
|
||
# 'source': '我是谁'
|
||
# }
|
||
# ]
|
||
#
|
||
# # 添加文档到知识库
|
||
# print("正在添加文档到知识库...")
|
||
# count = rag_system.add_documents(documents)
|
||
# print(f"成功添加 {count} 个文档块到知识库")
|
||
|
||
# 示例查询
|
||
questions = [
|
||
"reisasol宝宝亲亲",
|
||
"你是谁",
|
||
# "upsertMusic怎么用,不是upsertMusic01"
|
||
]
|
||
|
||
for question in questions:
|
||
print(f"\n问题: {question}")
|
||
result = rag_system.role_play_query(question,"Reisasol")
|
||
print(f"答案: {result['answer']}")
|
||
print("参考文档:")
|
||
for i, doc in enumerate(result['retrieved_docs'], 1):
|
||
print(f" {i}. {doc['text'][:100]}... (来源: {doc['source']})")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|