1

org/main.py

@@ -0,0 +1,38 @@
+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()