File size: 3,011 Bytes
4ce93da
 
b868160
3e48a1e
b868160
 
 
a2b8ed7
 
b868160
a2b8ed7
bc96608
4ce93da
 
6f31366
4ce93da
a2b8ed7
 
4ce93da
 
 
 
 
 
 
 
a2b8ed7
86e3d75
 
 
 
 
 
 
6f31366
86e3d75
 
 
 
6f31366
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
86e3d75
cd81c16
 
 
9aadbe4
6ab6544
cd81c16
 
 
9aadbe4
83f0301
9aadbe4
cd81c16
9aadbe4
 
 
cd81c16
 
 
9aadbe4
cd81c16
 
9aadbe4
6ab6544
cd81c16
9aadbe4
cd81c16
 
6ab6544
cd81c16
6ab6544
 
4ce93da
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
from flask import Flask, jsonify, request, Response, stream_with_context
from flask_cors import CORS
import requests
from bs4 import BeautifulSoup
import os
import re
import urllib.parse
import time
import random
import base64
from io import BytesIO
from googlesearch import search
import logging
import queue
from huggingface_hub import HfApi


app = Flask(__name__)
# Enable CORS with specific settings
CORS(app, resources={
    r"/*": {
        "origins": "*",
        "methods": ["GET", "POST", "OPTIONS"],
        "allow_headers": ["Content-Type", "Authorization"]
    }
})


HF_TOKEN = os.getenv("HF_TOKEN")  # Make sure you set the HF_TOKEN in your environment

@app.route('/restart_space', methods=['POST'])
def api_restart_space():
    """API route to restart a Hugging Face Space."""
    space_id = 'Pamudu13/web-scraper'  
    factory_reboot = request.json.get('factory_reboot', False)  # Optional: Set to True if you want a factory reboot

    if not space_id:
        return jsonify({'error': 'space_id parameter is required'}), 400

    try:
        hfapi = HfApi()

        # Call the restart_space method
        res = hfapi.restart_space(
            space_id,
            token=HF_TOKEN,
            factory_reboot=factory_reboot
        )

        return jsonify({
            'success': True,
            'message': f"Successfully restarted Space: {space_id}",
            'response': res
        }), 200

    except Exception as e:
        return jsonify({
            'success': False,
            'message': f"Error: {str(e)}"
        }), 500    

@app.route('/get_live_space_status', methods=['GET'])
def get_live_space_status():
    """API route to stream live status of a Hugging Face Space."""
    space_id = request.args.get('space_id', 'Pamudu13/web-scraper')  # Default to 'Pamudu13/web-scraper' if not provided

    def generate():
        while True:
            try:
                # Fetch the current runtime status of the Space
                hf_api = HfApi()
                space_runtime = hf_api.get_space_runtime(repo_id=space_id)
                
                # Extract relevant details
                status = space_runtime.stage  # e.g., 'BUILDING', 'RUNNING', etc.
                hardware = space_runtime.hardware  # e.g., 'cpu-basic', 't4-medium', etc.
                
                # Send the status as a Server-Sent Event
                yield f"data: {status}\n\n"
                yield f"data: {hardware}\n\n"

                # Delay before checking the status again
                time.sleep(5)  # Adjust polling interval as needed

            except Exception as e:
                # Handle errors and send an error message
                yield f"data: Error: {str(e)}\n\n"
                break  # Stop the stream in case of an error

    return Response(stream_with_context(generate()), mimetype='text/event-stream')


if __name__ == '__main__':
    logger.info("Starting Flask API server...")
    app.run(host='0.0.0.0', port=5001, debug=True)