1

org/app.py

@@ -0,0 +1,99 @@
+# file: app.py
+from flask import Flask, request, jsonify
+from rag_system import RAGSystem
+
+app = Flask(__name__)
+rag_system = RAGSystem()
+
+
+@app.route('/add_documents', methods=['POST'])
+def add_documents():
+    """
+    添加文档到知识库
+    请求体: JSON对象，包含documents字段
+    """
+    try:
+        data = request.get_json()
+        documents = data.get('documents', [])
+
+        if not documents:
+            return jsonify({'error': 'No documents provided'}), 400
+
+        count = rag_system.add_documents(documents)
+        return jsonify({
+            'message': f'Successfully added {count} document chunks to the knowledge base',
+            'added_count': count
+        }), 201
+
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+
+@app.route('/query', methods=['POST'])
+def query():
+    """
+    查询知识库
+    请求体: JSON对象，包含question字段和可选的top_k参数
+    """
+    try:
+        data = request.get_json()
+        question = data.get('question')
+        top_k = data.get('top_k', 3)
+
+        if not question:
+            return jsonify({'error': 'Question is required'}), 400
+
+        result = rag_system.query(question, top_k)
+        return jsonify(result), 200
+
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+
+@app.route('/add_songs', methods=['POST'])
+def add_songs():
+    """
+    添加歌曲数据到知识库
+    请求体: JSON对象，包含songs_data字段
+    """
+    try:
+        data = request.get_json()
+        songs_data = data.get('songs_data', [])
+
+        if not songs_data:
+            return jsonify({'error': 'No songs data provided'}), 400
+
+        count = rag_system.add_song_data(songs_data)
+        return jsonify({
+            'message': f'Successfully added {count} song documents to the knowledge base',
+            'added_count': count
+        }), 201
+
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+
+@app.route('/role_query', methods=['POST'])
+def role_query():
+    """
+    角色扮演查询
+    请求体: JSON对象，包含question, role字段和可选的top_k参数
+    """
+    try:
+        data = request.get_json()
+        question = data.get('question')
+        role = data.get('role')
+        top_k = data.get('top_k', 3)
+
+        if not question or not role:
+            return jsonify({'error': 'Question and role are required'}), 400
+
+        result = rag_system.role_play_query(question, role, top_k)
+        return jsonify(result), 200
+
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000, debug=True)