Update app.py

app.py

@@ -1,6 +1,6 @@
-from flask import Flask, render_template, request, jsonify
+from flask import Flask, render_template, request, jsonify, send_file
 from parser import parse_source_to_graph
-from dataset_gen import create_dataset_entry
+from dataset_gen import create_dataset_entry, get_dataset_stats, upload_to_hub, OUTPUT_FILE
 import os
 
 app = Flask(__name__)
@@ -11,24 +11,28 @@ def index():
 
 @app.route('/parse', methods=['POST'])
 def parse():
-    data = request.get_json()
-    code = data.get('code', '')
-    if not code.strip():
-        return jsonify({"nodes": [], "connections": []})
-        
-    result = parse_source_to_graph(code)
-    return jsonify(result)
-
-@app.route('/add_to_dataset', methods=['POST'])
-def add_to_dataset():
-    data = request.get_json()
-    code = data.get('code', '')
-    if not code.strip():
-        return jsonify({"status": "error", "message": "Empty code"})
-        
-    result = create_dataset_entry(code)
-    return jsonify(result)
+    code = request.json.get('code', '')
+    return jsonify(parse_source_to_graph(code))
+
+@app.route('/dataset/add', methods=['POST'])
+def add_entry():
+    code = request.json.get('code', '')
+    return jsonify(create_dataset_entry(code))
+
+@app.route('/dataset/list', methods=['GET'])
+def list_entries():
+    return jsonify(get_dataset_stats())
+
+@app.route('/dataset/download', methods=['GET'])
+def download_dataset():
+    if os.path.exists(OUTPUT_FILE):
+        return send_file(OUTPUT_FILE, as_attachment=True)
+    return "No dataset found", 404
+
+@app.route('/dataset/upload_hf', methods=['POST'])
+def upload_hf():
+    data = request.json
+    return jsonify(upload_to_hub(data.get('token'), data.get('repo_id')))
 
 if __name__ == '__main__':
-    # Running on 0.0.0.0 to be accessible if deployed or in containers
     app.run(host="0.0.0.0", port=7860, debug=True)