Update app.py

app.py

@@ -6,6 +6,7 @@ import json
 import random
 import uuid
 import concurrent.futures
+import threading
 from datetime import datetime, timedelta
 from apscheduler.schedulers.background import BackgroundScheduler
 from flask import Flask, request, jsonify, Response, stream_with_context
@@ -33,6 +34,10 @@ valid_keys_global = []
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=20)
 model_key_indices = {}
 
+request_timestamps = []
+token_counts = []
+data_lock = threading.Lock()
+
 def get_credit_summary(api_key):
     """
     使用 API 密钥获取额度信息。
@@ -394,7 +399,19 @@ scheduler.add_job(refresh_models, 'interval', hours=1)
 
 @app.route('/')
 def index():
-    return "<h1>Welcome to SiliconFlow</h1>"
+    current_time = time.time()
+    one_minute_ago = current_time - 60
+
+    with data_lock:
+        # Clean up old data
+        while request_timestamps and request_timestamps[0] < one_minute_ago:
+            request_timestamps.pop(0)
+            token_counts.pop(0)
+
+        rpm = len(request_timestamps)
+        tpm = sum(token_counts)
+
+    return jsonify({"rpm": rpm, "tpm": tpm})
 
 @app.route('/check_tokens', methods=['POST'])
 def check_tokens():
@@ -616,6 +633,10 @@ def handsome_chat_completions():
                     f"输出的内容: {response_content_replaced}"
                 )
 
+                with data_lock:
+                    request_timestamps.append(time.time())
+                    token_counts.append(prompt_tokens+completion_tokens)
+
             return Response(
                 stream_with_context(generate()),
                 content_type=response.headers['Content-Type']
@@ -678,6 +699,13 @@ def handsome_chat_completions():
                 f"用户的内容: {user_content_replaced}, "
                 f"输出的内容: {response_content_replaced}"
             )
+            with data_lock:
+                request_timestamps.append(time.time())
+                if "prompt_tokens" in response_json["usage"] and "completion_tokens" in response_json["usage"]:
+                    token_counts.append(response_json["usage"]["prompt_tokens"] + response_json["usage"]["completion_tokens"])
+                else:
+                    token_counts.append(0)
+
             return jsonify(response_json)
 
     except requests.exceptions.RequestException as e: