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()