125 lines
4.2 KiB
Python
125 lines
4.2 KiB
Python
#####################################
|
|
#
|
|
# If you meet 429 when running this script, please increase the `GLOBAL_API_RATE_LIMIT` in your Chat-API service.
|
|
# See more details in https://github.com/ai365vip/chat-api?tab=readme-ov-file#%E7%8E%AF%E5%A2%83%E5%8F%98%E9%87%8F
|
|
#
|
|
#####################################
|
|
|
|
import argparse
|
|
import concurrent.futures
|
|
|
|
from .oneapi_manager import OneAPIManager
|
|
from .cursor import Cursor
|
|
|
|
|
|
def handle_oneapi_cursor_channel(
|
|
oneapi: OneAPIManager,
|
|
channel_id,
|
|
test_channel: bool,
|
|
disable_low_balance_channel: bool,
|
|
delete_low_balance_channel: bool,
|
|
low_balance_threshold=10,
|
|
):
|
|
if test_channel:
|
|
test_response = oneapi.test_channel(channel_id)
|
|
if test_response.status_code != 200:
|
|
print(
|
|
f"Fail to test channel {channel_id}. Status Code: {response.status_code}"
|
|
)
|
|
return None
|
|
|
|
response = oneapi.get_channel(channel_id)
|
|
if response.status_code != 200:
|
|
print(f"Fail to get channel {channel_id}. Status Code: {response.status_code}")
|
|
return None
|
|
|
|
data = response.json()["data"]
|
|
key = data["key"]
|
|
status = data["status"] # 1 for enable, 2 for disbale
|
|
test_time = data["test_time"]
|
|
response_time = data["response_time"]
|
|
remaining_balance = Cursor.get_remaining_balance(key)
|
|
remaining_days = Cursor.get_trial_remaining_days(key)
|
|
print(
|
|
f"[OneAPI] Channel {channel_id} Info: Balance = {remaining_balance}. Trial Remaining Days = {remaining_days}. Response Time = {response_time}"
|
|
)
|
|
|
|
if None in [remaining_balance, remaining_days]:
|
|
print("[OneAPI] Invalid resposne")
|
|
return None
|
|
|
|
if remaining_balance < low_balance_threshold or (
|
|
test_time != 0 and response_time < 1000
|
|
): # or remaining_days <= 0:
|
|
if delete_low_balance_channel:
|
|
response = oneapi.delete_channel(channel_id)
|
|
print(
|
|
f"[OneAPI] Delete Channel {channel_id}. Status Coue: {response.status_code}"
|
|
)
|
|
elif disable_low_balance_channel and status == 1:
|
|
response = oneapi.disable_channel(channel_id)
|
|
print(
|
|
f"[OneAPI] Disable Channel {channel_id}. Status Code: {response.status_code}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--oneapi_url", type=str, required=False, help="URL link for One-API website"
|
|
)
|
|
parser.add_argument(
|
|
"--oneapi_token", type=str, required=False, help="Token for One-API website"
|
|
)
|
|
parser.add_argument(
|
|
"--test_channel", default=False, type=lambda x: (str(x).lower() == "true")
|
|
)
|
|
parser.add_argument(
|
|
"--disable_low_balance_accounts",
|
|
default=False,
|
|
type=lambda x: (str(x).lower() == "true"),
|
|
)
|
|
parser.add_argument(
|
|
"--delete_low_balance_accounts",
|
|
default=False,
|
|
type=lambda x: (str(x).lower() == "true"),
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--max_workers",
|
|
type=int,
|
|
default=10,
|
|
help="How many workers in multi-threading",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
oneapi_url = args.oneapi_url
|
|
oneapi_token = args.oneapi_token
|
|
test_channel = args.test_channel
|
|
disable_low_balance_accounts = args.disable_low_balance_accounts
|
|
delete_low_balance_accounts = args.delete_low_balance_accounts
|
|
max_workers = args.max_workers
|
|
|
|
oneapi = OneAPIManager(oneapi_url, oneapi_token)
|
|
|
|
response_channels = oneapi.get_channels(0, 2147483647)
|
|
channels = response_channels.json()["data"]
|
|
channels_ids = [channel["id"] for channel in channels]
|
|
channels_ids = sorted(channels_ids, key=int)
|
|
print(f"[OneAPI] Channel Count: {len(channels_ids)}")
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
futures = [
|
|
executor.submit(
|
|
handle_oneapi_cursor_channel,
|
|
oneapi,
|
|
id,
|
|
test_channel,
|
|
disable_low_balance_accounts,
|
|
delete_low_balance_accounts,
|
|
)
|
|
for id in channels_ids
|
|
]
|
|
for future in concurrent.futures.as_completed(futures):
|
|
result = future.result()
|