1

org/add.py

@@ -0,0 +1,61 @@
+import json
+from rag_system import RAGSystem
+
+
+def load_songs_from_json(file_path):
+    """从JSON文件加载歌曲数据"""
+    try:
+        with open(file_path, 'r', encoding='utf-8') as f:
+            songs_data = json.load(f)
+        return songs_data
+    except FileNotFoundError:
+        print(f"文件 {file_path} 未找到")
+        return []
+    except json.JSONDecodeError as e:
+        print(f"JSON解析错误: {e}")
+        return []
+
+
+def convert_songs_to_documents(songs_data):
+    """将歌曲数据转换为文档格式"""
+    documents = []
+    for song in songs_data:
+        content = f"歌曲名称: {song.get('title', '未知')}, 歌手: {song.get('artist', '未知')}, BPM: {song.get('bpm', '未知')}, 版本: {song.get('version', '未知')}"
+        documents.append({
+            'content': content,
+            'source': f"歌曲数据 - {song.get('title', '未知歌曲')}"
+        })
+    return documents
+
+
+def main():
+    # 初始化RAG系统
+    rag_system = RAGSystem()
+
+    # 从JSON文件读取歌曲数据并添加到知识库
+    songs_file = "./put.json"
+    songs_data = load_songs_from_json(songs_file)
+    if songs_data:
+        print("正在添加歌曲数据到知识库...")
+        song_documents = convert_songs_to_documents(songs_data)
+        count = rag_system.add_documents(song_documents)
+        print(f"成功添加 {count} 个歌曲文档到知识库")
+
+    # 示例查询
+    questions = [
+        "pandora怎么样",
+        "你是谁",
+        # "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()