Spaces:
Running
Running
from flask import Flask, request, jsonify, Response, render_template_string, render_template, redirect, url_for, session as flask_session | |
import requests | |
import time | |
import json | |
import uuid | |
import random | |
import io | |
import re | |
from functools import wraps | |
import hashlib | |
import jwt | |
import os | |
import threading | |
from datetime import datetime, timedelta | |
app = Flask(__name__, template_folder='templates') | |
app.secret_key = os.environ.get("SECRET_KEY", "abacus_chat_proxy_secret_key") | |
app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(days=7) | |
# 添加tokenizer服务URL | |
TOKENIZER_SERVICE_URL = "https://malt666-tokenizer.hf.space/count_tokens" | |
API_ENDPOINT_URL = "https://abacus.ai/api/v0/describeDeployment" | |
MODEL_LIST_URL = "https://abacus.ai/api/v0/listExternalApplications" | |
CHAT_URL = "https://apps.abacus.ai/api/_chatLLMSendMessageSSE" | |
USER_INFO_URL = "https://abacus.ai/api/v0/_getUserInfo" | |
COMPUTE_POINTS_URL = "https://apps.abacus.ai/api/_getOrganizationComputePoints" | |
COMPUTE_POINTS_LOG_URL = "https://abacus.ai/api/v0/_getOrganizationComputePointLog" | |
CREATE_CONVERSATION_URL = "https://apps.abacus.ai/api/createDeploymentConversation" | |
DELETE_CONVERSATION_URL = "https://apps.abacus.ai/api/deleteDeploymentConversation" | |
GET_CONVERSATION_URL = "https://apps.abacus.ai/api/getDeploymentConversation" | |
COMPUTE_POINT_TOGGLE_URL = "https://abacus.ai/api/v0/_updateOrganizationComputePointToggle" | |
USER_AGENTS = [ | |
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/116.0.0.0 Safari/537.36" | |
] | |
PASSWORD = None | |
USER_NUM = 0 | |
USER_DATA = [] | |
CURRENT_USER = -1 | |
MODELS = set() | |
# 添加线程锁用于保护 CURRENT_USER 的访问 | |
user_selection_lock = threading.Lock() | |
TRACE_ID = "3042e28b3abf475d8d973c7e904935af" | |
SENTRY_TRACE = f"{TRACE_ID}-80d9d2538b2682d0" | |
# 添加一个计数器记录健康检查次数 | |
health_check_counter = 0 | |
# 添加统计变量 | |
model_usage_stats = {} # 模型使用次数统计 | |
total_tokens = { | |
"prompt": 0, # 输入token统计 | |
"completion": 0, # 输出token统计 | |
"total": 0 # 总token统计 | |
} | |
# 模型调用记录 | |
model_usage_records = [] # 每次调用详细记录 | |
MODEL_USAGE_RECORDS_FILE = "model_usage_records.json" # 调用记录保存文件 | |
# 计算点信息 | |
compute_points = { | |
"left": 0, # 剩余计算点 | |
"total": 0, # 总计算点 | |
"used": 0, # 已使用计算点 | |
"percentage": 0, # 使用百分比 | |
"last_update": None # 最后更新时间 | |
} | |
# 计算点使用日志 | |
compute_points_log = { | |
"columns": {}, # 列名 | |
"log": [] # 日志数据 | |
} | |
# 多用户计算点信息 | |
users_compute_points = [] | |
# 记录启动时间 | |
START_TIME = datetime.utcnow() + timedelta(hours=8) # 北京时间 | |
# 自定义JSON编码器,处理datetime对象 | |
class DateTimeEncoder(json.JSONEncoder): | |
def default(self, obj): | |
if isinstance(obj, datetime): | |
return obj.strftime('%Y-%m-%d %H:%M:%S') | |
return super(DateTimeEncoder, self).default(obj) | |
# 加载模型调用记录 | |
def load_model_usage_records(): | |
global model_usage_records | |
try: | |
if os.path.exists(MODEL_USAGE_RECORDS_FILE): | |
with open(MODEL_USAGE_RECORDS_FILE, 'r', encoding='utf-8') as f: | |
records = json.load(f) | |
if isinstance(records, list): | |
model_usage_records = records | |
print(f"成功加载 {len(model_usage_records)} 条模型调用记录") | |
else: | |
print("调用记录文件格式不正确,初始化为空列表") | |
except Exception as e: | |
print(f"加载模型调用记录失败: {e}") | |
model_usage_records = [] | |
# 保存模型调用记录 | |
def save_model_usage_records(): | |
try: | |
with open(MODEL_USAGE_RECORDS_FILE, 'w', encoding='utf-8') as f: | |
json.dump(model_usage_records, f, ensure_ascii=False, indent=2, cls=DateTimeEncoder) | |
print(f"成功保存 {len(model_usage_records)} 条模型调用记录") | |
except Exception as e: | |
print(f"保存模型调用记录失败: {e}") | |
def update_conversation_id(user_index, conversation_id): | |
"""更新用户的conversation_id并保存到配置文件""" | |
try: | |
with open("config.json", "r") as f: | |
config = json.load(f) | |
if "config" in config and user_index < len(config["config"]): | |
config["config"][user_index]["conversation_id"] = conversation_id | |
# 保存到配置文件 | |
with open("config.json", "w") as f: | |
json.dump(config, f, indent=4) | |
print(f"已将用户 {user_index+1} 的conversation_id更新为: {conversation_id}") | |
else: | |
print(f"更新conversation_id失败: 配置文件格式错误或用户索引越界") | |
except Exception as e: | |
print(f"更新conversation_id失败: {e}") | |
def resolve_config(): | |
# 从环境变量读取多组配置 | |
config_list = [] | |
i = 1 | |
while True: | |
cookie = os.environ.get(f"cookie_{i}") | |
if not cookie: | |
break | |
# 为每个cookie创建一个配置项,conversation_id初始为空 | |
config_list.append({ | |
"conversation_id": "", # 初始为空,将通过get_or_create_conversation自动创建 | |
"cookies": cookie | |
}) | |
i += 1 | |
# 如果环境变量存在配置,使用环境变量的配置 | |
if config_list: | |
print(f"从环境变量加载了 {len(config_list)} 个配置") | |
return config_list | |
# 如果环境变量不存在,从文件读取 | |
try: | |
with open("config.json", "r") as f: | |
config = json.load(f) | |
config_list = config.get("config") | |
return config_list | |
except FileNotFoundError: | |
print("未找到config.json文件") | |
return [] | |
except json.JSONDecodeError: | |
print("config.json格式错误") | |
return [] | |
def get_password(): | |
global PASSWORD | |
# 从环境变量读取密码 | |
env_password = os.environ.get("password") | |
if env_password: | |
PASSWORD = hashlib.sha256(env_password.encode()).hexdigest() | |
return | |
# 如果环境变量不存在,从文件读取 | |
try: | |
with open("password.txt", "r") as f: | |
PASSWORD = f.read().strip() | |
except FileNotFoundError: | |
with open("password.txt", "w") as f: | |
PASSWORD = None | |
def require_auth(f): | |
def decorated(*args, **kwargs): | |
if not PASSWORD: | |
return f(*args, **kwargs) | |
# 检查Flask会话是否已登录 | |
if flask_session.get('logged_in'): | |
return f(*args, **kwargs) | |
# 如果是API请求,检查Authorization头 | |
auth = request.authorization | |
if not auth or not check_auth(auth.token): | |
# 如果是浏览器请求,重定向到登录页面 | |
if request.headers.get('Accept', '').find('text/html') >= 0: | |
return redirect(url_for('login')) | |
return jsonify({"error": "Unauthorized access"}), 401 | |
return f(*args, **kwargs) | |
return decorated | |
def check_auth(token): | |
return hashlib.sha256(token.encode()).hexdigest() == PASSWORD | |
def is_token_expired(token): | |
if not token: | |
return True | |
try: | |
# Malkodi tokenon sen validigo de subskribo | |
payload = jwt.decode(token, options={"verify_signature": False}) | |
# Akiru eksvalidiĝan tempon, konsideru eksvalidiĝinta 5 minutojn antaŭe | |
return payload.get('exp', 0) - time.time() < 300 | |
except: | |
return True | |
def refresh_token(session, cookies): | |
"""Uzu kuketon por refreŝigi session token, nur revenigu novan tokenon""" | |
headers = { | |
"accept": "application/json, text/plain, */*", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"content-type": "application/json", | |
"reai-ui": "1", | |
"sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"", | |
"sec-ch-ua-mobile": "?0", | |
"sec-ch-ua-platform": "\"Windows\"", | |
"sec-fetch-dest": "empty", | |
"sec-fetch-mode": "cors", | |
"sec-fetch-site": "same-site", | |
"x-abacus-org-host": "apps", | |
"user-agent": random.choice(USER_AGENTS), | |
"origin": "https://apps.abacus.ai", | |
"referer": "https://apps.abacus.ai/", | |
"cookie": cookies | |
} | |
try: | |
response = session.post( | |
USER_INFO_URL, | |
headers=headers, | |
json={}, | |
cookies=None | |
) | |
if response.status_code == 200: | |
response_data = response.json() | |
if response_data.get('success') and 'sessionToken' in response_data.get('result', {}): | |
return response_data['result']['sessionToken'] | |
else: | |
print(f"刷新token失败: {response_data.get('error', '未知错误')}") | |
return None | |
else: | |
print(f"刷新token失败,状态码: {response.status_code}") | |
return None | |
except Exception as e: | |
print(f"刷新token异常: {e}") | |
return None | |
def get_model_map(session, cookies, session_token): | |
"""Akiru disponeblan modelan liston kaj ĝiajn mapajn rilatojn""" | |
headers = { | |
"accept": "application/json, text/plain, */*", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"content-type": "application/json", | |
"reai-ui": "1", | |
"sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"", | |
"sec-ch-ua-mobile": "?0", | |
"sec-ch-ua-platform": "\"Windows\"", | |
"sec-fetch-dest": "empty", | |
"sec-fetch-mode": "cors", | |
"sec-fetch-site": "same-site", | |
"x-abacus-org-host": "apps", | |
"user-agent": random.choice(USER_AGENTS), | |
"origin": "https://apps.abacus.ai", | |
"referer": "https://apps.abacus.ai/", | |
"cookie": cookies | |
} | |
if session_token: | |
headers["session-token"] = session_token | |
model_map = {} | |
models_set = set() | |
try: | |
response = session.post( | |
MODEL_LIST_URL, | |
headers=headers, | |
json={}, | |
cookies=None | |
) | |
if response.status_code != 200: | |
print(f"获取模型列表失败,状态码: {response.status_code}") | |
raise Exception("API请求失败") | |
data = response.json() | |
if not data.get('success'): | |
print(f"获取模型列表失败: {data.get('error', '未知错误')}") | |
raise Exception("API返回错误") | |
applications = [] | |
if isinstance(data.get('result'), dict): | |
applications = data.get('result', {}).get('externalApplications', []) | |
elif isinstance(data.get('result'), list): | |
applications = data.get('result', []) | |
for app in applications: | |
app_name = app.get('name', '') | |
app_id = app.get('externalApplicationId', '') | |
prediction_overrides = app.get('predictionOverrides', {}) | |
llm_name = prediction_overrides.get('llmName', '') if prediction_overrides else '' | |
if not (app_name and app_id and llm_name): | |
continue | |
model_name = app_name | |
model_map[model_name] = (app_id, llm_name) | |
models_set.add(model_name) | |
if not model_map: | |
raise Exception("未找到任何可用模型") | |
return model_map, models_set | |
except Exception as e: | |
print(f"获取模型列表异常: {e}") | |
raise | |
def init_session(): | |
get_password() | |
global USER_NUM, MODELS, USER_DATA | |
config_list = resolve_config() | |
user_num = len(config_list) | |
all_models = set() | |
for i in range(user_num): | |
user = config_list[i] | |
cookies = user.get("cookies") | |
conversation_id = user.get("conversation_id") | |
session = requests.Session() | |
session_token = refresh_token(session, cookies) | |
if not session_token: | |
print(f"无法获取cookie {i+1}的token") | |
continue | |
try: | |
model_map, models_set = get_model_map(session, cookies, session_token) | |
all_models.update(models_set) | |
USER_DATA.append((session, cookies, session_token, conversation_id, model_map, i)) | |
# 对第一个成功配置的用户,初始化计算点数记录功能 | |
if i == 0: | |
try: | |
headers = { | |
"accept": "application/json, text/plain, */*", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"content-type": "application/json", | |
"reai-ui": "1", | |
"sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"", | |
"sec-ch-ua-mobile": "?0", | |
"sec-ch-ua-platform": "\"Windows\"", | |
"sec-fetch-dest": "empty", | |
"sec-fetch-mode": "cors", | |
"sec-fetch-site": "same-site", | |
"x-abacus-org-host": "apps", | |
"session-token": session_token | |
} | |
response = session.post( | |
COMPUTE_POINT_TOGGLE_URL, | |
headers=headers, | |
json={"alwaysDisplay": True}, | |
cookies=None | |
) | |
if response.status_code == 200: | |
result = response.json() | |
if result.get("success"): | |
print("成功初始化计算点数记录功能为开启状态") | |
else: | |
print(f"初始化计算点数记录功能失败: {result.get('error', '未知错误')}") | |
else: | |
print(f"初始化计算点数记录功能失败,状态码: {response.status_code}") | |
except Exception as e: | |
print(f"初始化计算点数记录功能时出错: {e}") | |
except Exception as e: | |
print(f"配置用户 {i+1} 失败: {e}") | |
continue | |
USER_NUM = len(USER_DATA) | |
if USER_NUM == 0: | |
print("No user available, exiting...") | |
exit(1) | |
MODELS = all_models | |
print(f"启动完成,共配置 {USER_NUM} 个用户") | |
def update_cookie(session, cookies): | |
cookie_jar = {} | |
for key, value in session.cookies.items(): | |
cookie_jar[key] = value | |
cookie_dict = {} | |
for item in cookies.split(";"): | |
key, value = item.strip().split("=", 1) | |
cookie_dict[key] = value | |
cookie_dict.update(cookie_jar) | |
cookies = "; ".join([f"{key}={value}" for key, value in cookie_dict.items()]) | |
return cookies | |
user_data = init_session() | |
def get_models(): | |
if len(MODELS) == 0: | |
return jsonify({"error": "No models available"}), 500 | |
model_list = [] | |
for model in MODELS: | |
model_list.append( | |
{ | |
"id": model, | |
"object": "model", | |
"created": int(time.time()), | |
"owned_by": "Elbert", | |
"name": model, | |
} | |
) | |
return jsonify({"object": "list", "data": model_list}) | |
def chat_completions(): | |
openai_request = request.get_json() | |
stream = openai_request.get("stream", False) | |
messages = openai_request.get("messages") | |
if messages is None: | |
return jsonify({"error": "Messages is required", "status": 400}), 400 | |
model = openai_request.get("model") | |
if model not in MODELS: | |
return ( | |
jsonify( | |
{ | |
"error": "Model not available, check if it is configured properly", | |
"status": 404, | |
} | |
), | |
404, | |
) | |
message = format_message(messages) | |
think = ( | |
openai_request.get("think", False) if model == "Claude Sonnet 3.7" else False | |
) | |
return ( | |
send_message(message, model, think) | |
if stream | |
else send_message_non_stream(message, model, think) | |
) | |
def get_user_data(): | |
global CURRENT_USER | |
# 使用锁确保线程安全 | |
with user_selection_lock: | |
CURRENT_USER = (CURRENT_USER + 1) % USER_NUM | |
current_user_index_local = CURRENT_USER # 本地副本,避免锁外访问全局变量 | |
print(f"使用配置 {current_user_index_local+1}") | |
# Akiru uzantajn datumojn (使用本地索引) | |
session, cookies, session_token, conversation_id, model_map, user_index = USER_DATA[current_user_index_local] | |
# Kontrolu ĉu la tokeno eksvalidiĝis, se jes, refreŝigu ĝin | |
if is_token_expired(session_token): | |
print(f"Cookie {current_user_index_local+1}的token已过期或即将过期,正在刷新...") | |
new_token = refresh_token(session, cookies) | |
if new_token: | |
# Ĝisdatigu la globale konservitan tokenon (加锁保护写入) | |
with user_selection_lock: | |
# 重新获取最新的 USER_DATA 状态再更新 | |
_session, _cookies, _, _conv_id, _model_map, _user_idx = USER_DATA[current_user_index_local] | |
USER_DATA[current_user_index_local] = (_session, _cookies, new_token, _conv_id, _model_map, _user_idx) | |
session_token = new_token # 更新函数内部使用的token | |
print(f"成功更新token: {session_token[:15]}...{session_token[-15:]}") | |
else: | |
print(f"警告:无法刷新Cookie {current_user_index_local+1}的token,继续使用当前token") | |
# 返回获取到的数据 (使用本地索引) | |
return (session, cookies, session_token, conversation_id, model_map, user_index) | |
def create_conversation(session, cookies, session_token, external_application_id=None, deployment_id=None): | |
"""创建新的会话""" | |
if not (external_application_id and deployment_id): | |
print("无法创建新会话: 缺少必要参数") | |
return None | |
headers = { | |
"accept": "application/json, text/plain, */*", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"content-type": "application/json", | |
"cookie": cookies, | |
"user-agent": random.choice(USER_AGENTS), | |
"x-abacus-org-host": "apps" | |
} | |
if session_token: | |
headers["session-token"] = session_token | |
create_payload = { | |
"deploymentId": deployment_id, | |
"name": "New Chat", | |
"externalApplicationId": external_application_id | |
} | |
try: | |
response = session.post( | |
CREATE_CONVERSATION_URL, | |
headers=headers, | |
json=create_payload | |
) | |
if response.status_code == 200: | |
data = response.json() | |
if data.get("success", False): | |
new_conversation_id = data.get("result", {}).get("deploymentConversationId") | |
if new_conversation_id: | |
print(f"成功创建新的conversation: {new_conversation_id}") | |
return new_conversation_id | |
print(f"创建会话失败: {response.status_code} - {response.text[:100]}") | |
return None | |
except Exception as e: | |
print(f"创建会话时出错: {e}") | |
return None | |
def delete_conversation(session, cookies, session_token, conversation_id, deployment_id="14b2a314cc"): | |
"""删除指定的对话""" | |
if not conversation_id: | |
print("无法删除对话: 缺少conversation_id") | |
return False | |
headers = { | |
"accept": "application/json, text/plain, */*", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"content-type": "application/json", | |
"cookie": cookies, | |
"user-agent": random.choice(USER_AGENTS), | |
"x-abacus-org-host": "apps" | |
} | |
if session_token: | |
headers["session-token"] = session_token | |
delete_payload = { | |
"deploymentId": deployment_id, | |
"deploymentConversationId": conversation_id | |
} | |
try: | |
response = session.post( | |
DELETE_CONVERSATION_URL, | |
headers=headers, | |
json=delete_payload | |
) | |
if response.status_code == 200: | |
data = response.json() | |
if data.get("success", False): | |
print(f"成功删除对话: {conversation_id}") | |
return True | |
print(f"删除对话失败: {response.status_code} - {response.text[:100]}") | |
return False | |
except Exception as e: | |
print(f"删除对话时出错: {e}") | |
return False | |
def get_or_create_conversation(session, cookies, session_token, conversation_id, model_map, model, user_index): | |
"""获取对话ID,如果不存在则创建;返回是否是使用现有会话""" | |
print(f"\n----- 获取会话ID (用户 {user_index+1}) -----") | |
# 如果有现有的会话ID,直接使用 | |
if conversation_id: | |
print(f"使用现有会话ID: {conversation_id}") | |
return conversation_id, True | |
# 如果没有会话ID,创建新的 | |
print("没有会话ID,创建新会话...") | |
deployment_id = "14b2a314cc" | |
# 确保模型信息存在 | |
if model not in model_map or len(model_map[model]) < 2: | |
print(f"错误: 无法获取模型 {model} 的信息") | |
return None, False | |
external_app_id = model_map[model][0] | |
# 创建新会话 | |
new_conversation_id = create_conversation( | |
session, cookies, session_token, | |
external_application_id=external_app_id, | |
deployment_id=deployment_id | |
) | |
if new_conversation_id: | |
print(f"成功创建新会话ID: {new_conversation_id}") | |
# 更新全局存储的会话ID | |
global USER_DATA, CURRENT_USER | |
session, cookies, session_token, _, model_map, _ = USER_DATA[CURRENT_USER] | |
USER_DATA[CURRENT_USER] = (session, cookies, session_token, new_conversation_id, model_map, user_index) | |
# 保存到配置文件 | |
update_conversation_id(user_index, new_conversation_id) | |
return new_conversation_id, False | |
print("创建新会话失败") | |
return None, False | |
def generate_trace_id(): | |
"""Generu novan trace_id kaj sentry_trace""" | |
trace_id = str(uuid.uuid4()).replace('-', '') | |
sentry_trace = f"{trace_id}-{str(uuid.uuid4())[:16]}" | |
return trace_id, sentry_trace | |
def send_message(message, model, think=False): | |
"""Flua traktado kaj plusendo de mesaĝoj""" | |
print("\n===== 开始处理消息 =====") | |
print(f"模型: {model}") | |
print(f"思考模式: {think}") | |
(session, cookies, session_token, conversation_id, model_map, user_index) = get_user_data() | |
print(f"使用用户配置: {user_index + 1}") | |
# 获取会话ID,并判断是否使用现有会话 | |
conversation_id, is_existing = get_or_create_conversation( | |
session, cookies, session_token, conversation_id, model_map, model, user_index | |
) | |
# 如果没有有效的会话ID,返回错误 | |
if not conversation_id: | |
return jsonify({"error": "Failed to get a valid conversation ID"}), 500 | |
print(f"会话ID: {conversation_id} (是否为现有会话: {is_existing})") | |
trace_id, sentry_trace = generate_trace_id() | |
# 计算输入token | |
completion_buffer = io.StringIO() # 收集所有输出用于计算token | |
headers = { | |
"accept": "text/event-stream", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"baggage": f"sentry-environment=production,sentry-release=975eec6685013679c139fc88db2c48e123d5c604,sentry-public_key=3476ea6df1585dd10e92cdae3a66ff49,sentry-trace_id={trace_id}", | |
"content-type": "text/plain;charset=UTF-8", | |
"cookie": cookies, | |
"sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"", | |
"sec-ch-ua-mobile": "?0", | |
"sec-ch-ua-platform": "\"Windows\"", | |
"sec-fetch-dest": "empty", | |
"sec-fetch-mode": "cors", | |
"sec-fetch-site": "same-origin", | |
"sentry-trace": sentry_trace, | |
"user-agent": random.choice(USER_AGENTS), | |
"x-abacus-org-host": "apps" | |
} | |
if session_token: | |
headers["session-token"] = session_token | |
# 构建基础请求 | |
payload = { | |
"requestId": str(uuid.uuid4()), | |
"deploymentConversationId": conversation_id, | |
"message": message, | |
"isDesktop": False, | |
"chatConfig": { | |
"timezone": "Asia/Shanghai", | |
"language": "zh-CN" | |
}, | |
"llmName": model_map[model][1], | |
"externalApplicationId": model_map[model][0] | |
} | |
# 如果是使用现有会话,添加regenerate和editPrompt参数 | |
if is_existing: | |
payload["regenerate"] = True | |
payload["editPrompt"] = True | |
print("为现有会话添加 regenerate=True 和 editPrompt=True") | |
if think: | |
payload["useThinking"] = think | |
try: | |
response = session.post( | |
CHAT_URL, | |
headers=headers, | |
data=json.dumps(payload), | |
stream=True, | |
cookies=None | |
) | |
response.raise_for_status() | |
def extract_segment(line_data): | |
try: | |
data = json.loads(line_data) | |
if "segment" in data: | |
if isinstance(data["segment"], str): | |
return data["segment"] | |
elif isinstance(data["segment"], dict) and "segment" in data["segment"]: | |
return data["segment"]["segment"] | |
return "" | |
except: | |
return "" | |
def generate(): | |
id = "" | |
think_state = 2 | |
yield "data: " + json.dumps({"object": "chat.completion.chunk", "choices": [{"delta": {"role": "assistant"}}]}) + "\n\n" | |
for line in response.iter_lines(): | |
if line: | |
decoded_line = line.decode("utf-8") | |
try: | |
if think: | |
data = json.loads(decoded_line) | |
if data.get("type") != "text": | |
continue | |
elif think_state == 2: | |
id = data.get("messageId") | |
segment = "<think>\n" + data.get("segment", "") | |
completion_buffer.write(segment) # 收集输出 | |
yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n" | |
think_state = 1 | |
elif think_state == 1: | |
if data.get("messageId") != id: | |
segment = data.get("segment", "") | |
completion_buffer.write(segment) | |
yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n" | |
else: | |
segment = "\n</think>\n" + data.get("segment", "") | |
completion_buffer.write(segment) | |
yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n" | |
think_state = 0 | |
else: | |
segment = data.get("segment", "") | |
completion_buffer.write(segment) | |
yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n" | |
else: | |
segment = extract_segment(decoded_line) | |
if segment: | |
completion_buffer.write(segment) | |
yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n" | |
except Exception as e: | |
print(f"处理响应出错: {e}") | |
# ---- Token 计算移到这里 ---- | |
print("\n----- 计算 Tokens (流式结束后) -----") | |
prompt_tokens, calculation_method = num_tokens_from_string(message, model) | |
print(f"输入 token 数: {prompt_tokens} (方法: {calculation_method})") | |
completion_content = completion_buffer.getvalue() | |
completion_tokens, comp_calc_method = num_tokens_from_string(completion_content, model) | |
print(f"输出 token 数: {completion_tokens} (方法: {comp_calc_method})") | |
# 决定最终使用的计算方法 (优先使用精确) | |
final_calculation_method = "精确" if calculation_method == "精确" and comp_calc_method == "精确" else "估算" | |
# ---- Token 计算结束 ---- | |
yield "data: " + json.dumps({"object": "chat.completion.chunk", "choices": [{"delta": {}, "finish_reason": "stop"}]}) + "\n\n" | |
yield "data: [DONE]\n\n" | |
# 在流式传输完成后计算token并更新统计 | |
# 注意: 如果客户端在流结束前断开连接,这里的 completion_content 可能不完整, | |
# 导致 completion_tokens 和 total_tokens 的本地记录不准确。 | |
# 但 Abacus 的计算点数扣除通常在其服务端完成,不受此影响。 | |
# 保存对话历史并获取计算点数 | |
_, compute_points_used = save_conversation_history(session, cookies, session_token, conversation_id) | |
# 更新统计信息 | |
update_model_stats(model, prompt_tokens, completion_tokens, final_calculation_method, compute_points_used) | |
return Response(generate(), mimetype="text/event-stream") | |
except requests.exceptions.RequestException as e: | |
error_details = str(e) | |
if hasattr(e, 'response') and e.response is not None: | |
if hasattr(e.response, 'text'): | |
error_details += f" - Response: {e.response.text[:200]}" | |
print(f"发送消息失败: {error_details}") | |
# 如果是使用现有会话失败,尝试创建新会话重试一次 | |
if is_existing: | |
print("使用现有会话失败,尝试创建新会话...") | |
# 创建新会话 | |
deployment_id = "14b2a314cc" | |
external_app_id = model_map[model][0] if model in model_map and len(model_map[model]) >= 2 else None | |
if external_app_id: | |
new_conversation_id = create_conversation( | |
session, cookies, session_token, | |
external_application_id=external_app_id, | |
deployment_id=deployment_id | |
) | |
if new_conversation_id: | |
print(f"成功创建新会话ID: {new_conversation_id},重试请求") | |
# 更新全局存储的会话ID | |
global USER_DATA, CURRENT_USER | |
session, cookies, session_token, _, model_map, _ = USER_DATA[CURRENT_USER] | |
USER_DATA[CURRENT_USER] = (session, cookies, session_token, new_conversation_id, model_map, user_index) | |
# 保存到配置文件 | |
update_conversation_id(user_index, new_conversation_id) | |
# 修改payload使用新会话ID,并移除regenerate和editPrompt | |
payload["deploymentConversationId"] = new_conversation_id | |
if "regenerate" in payload: | |
del payload["regenerate"] | |
if "editPrompt" in payload: | |
del payload["editPrompt"] | |
try: | |
# 非流式重试逻辑与流式类似,但需要重新提取响应内容 | |
response = session.post( | |
CHAT_URL, | |
headers=headers, | |
data=json.dumps(payload), | |
stream=True, | |
cookies=None | |
) | |
response.raise_for_status() | |
# 重用现有提取逻辑... | |
# 但这里代码重复太多,实际应该重构为共享函数 | |
buffer = io.StringIO() | |
for line in response.iter_lines(): | |
if line: | |
decoded_line = line.decode("utf-8") | |
segment = extract_segment(decoded_line) | |
if segment: | |
buffer.write(segment) | |
response_content = buffer.getvalue() | |
# ---- 重试逻辑中的 Token 计算 ---- | |
print("\n----- 计算 Tokens (重试成功后) -----") | |
prompt_tokens, calculation_method = num_tokens_from_string(message, model) | |
print(f"输入 token 数: {prompt_tokens} (方法: {calculation_method})") | |
# 计算输出token并更新统计信息 | |
completion_tokens, comp_calc_method = num_tokens_from_string(response_content, model) | |
print(f"输出 token 数: {completion_tokens} (方法: {comp_calc_method})") | |
# 决定最终使用的计算方法 | |
final_calculation_method = "精确" if calculation_method == "精确" and comp_calc_method == "精确" else "估算" | |
# ---- Token 计算结束 ---- | |
# 保存对话历史并获取计算点数 | |
_, compute_points_used = save_conversation_history(session, cookies, session_token, new_conversation_id) | |
# 更新统计信息 | |
# 注意: 重试逻辑。Token 计算准确性依赖于 response 完整性。 | |
update_model_stats(model, prompt_tokens, completion_tokens, final_calculation_method, compute_points_used) | |
return jsonify({ | |
"id": f"chatcmpl-{str(uuid.uuid4())}", | |
"object": "chat.completion", | |
"created": int(time.time()), | |
"model": model, | |
"choices": [{ | |
"index": 0, | |
"message": { | |
"role": "assistant", | |
"content": response_content | |
}, | |
"finish_reason": "stop" | |
}], | |
"usage": { | |
"prompt_tokens": prompt_tokens, | |
"completion_tokens": completion_tokens, | |
"total_tokens": prompt_tokens + completion_tokens | |
} | |
}) | |
except Exception as retry_e: | |
print(f"重试失败: {retry_e}") | |
return jsonify({"error": f"Failed to send message: {error_details}"}), 500 | |
def send_message_non_stream(message, model, think=False): | |
"""Ne-flua traktado de mesaĝoj""" | |
print("\n===== 开始处理消息(非流式) =====") | |
print(f"模型: {model}") | |
print(f"思考模式: {think}") | |
(session, cookies, session_token, conversation_id, model_map, user_index) = get_user_data() | |
print(f"使用用户配置: {user_index + 1}") | |
# 获取会话ID,并判断是否使用现有会话 | |
conversation_id, is_existing = get_or_create_conversation( | |
session, cookies, session_token, conversation_id, model_map, model, user_index | |
) | |
# 如果没有有效的会话ID,返回错误 | |
if not conversation_id: | |
return jsonify({"error": "Failed to get a valid conversation ID"}), 500 | |
print(f"会话ID: {conversation_id} (是否为现有会话: {is_existing})") | |
trace_id, sentry_trace = generate_trace_id() | |
headers = { | |
"accept": "text/event-stream", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"baggage": f"sentry-environment=production,sentry-release=975eec6685013679c139fc88db2c48e123d5c604,sentry-public_key=3476ea6df1585dd10e92cdae3a66ff49,sentry-trace_id={trace_id}", | |
"content-type": "text/plain;charset=UTF-8", | |
"cookie": cookies, | |
"sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"", | |
"sec-ch-ua-mobile": "?0", | |
"sec-ch-ua-platform": "\"Windows\"", | |
"sec-fetch-dest": "empty", | |
"sec-fetch-mode": "cors", | |
"sec-fetch-site": "same-origin", | |
"sentry-trace": sentry_trace, | |
"user-agent": random.choice(USER_AGENTS), | |
"x-abacus-org-host": "apps" | |
} | |
if session_token: | |
headers["session-token"] = session_token | |
# 构建基础请求 | |
payload = { | |
"requestId": str(uuid.uuid4()), | |
"deploymentConversationId": conversation_id, | |
"message": message, | |
"isDesktop": False, | |
"chatConfig": { | |
"timezone": "Asia/Shanghai", | |
"language": "zh-CN" | |
}, | |
"llmName": model_map[model][1], | |
"externalApplicationId": model_map[model][0] | |
} | |
# 如果是使用现有会话,添加regenerate和editPrompt参数 | |
if is_existing: | |
payload["regenerate"] = True | |
payload["editPrompt"] = True | |
print("为现有会话添加 regenerate=True 和 editPrompt=True") | |
if think: | |
payload["useThinking"] = think | |
try: | |
response = session.post( | |
CHAT_URL, | |
headers=headers, | |
data=json.dumps(payload), | |
stream=True, | |
cookies=None | |
) | |
response.raise_for_status() | |
buffer = io.StringIO() | |
def extract_segment(line_data): | |
try: | |
data = json.loads(line_data) | |
if "segment" in data: | |
if isinstance(data["segment"], str): | |
return data["segment"] | |
elif isinstance(data["segment"], dict) and "segment" in data["segment"]: | |
return data["segment"]["segment"] | |
return "" | |
except: | |
return "" | |
if think: | |
id = "" | |
think_state = 2 | |
think_buffer = io.StringIO() | |
content_buffer = io.StringIO() | |
for line in response.iter_lines(): | |
if line: | |
decoded_line = line.decode("utf-8") | |
try: | |
data = json.loads(decoded_line) | |
if data.get("type") != "text": | |
continue | |
elif think_state == 2: | |
id = data.get("messageId") | |
segment = data.get("segment", "") | |
think_buffer.write(segment) | |
think_state = 1 | |
elif think_state == 1: | |
if data.get("messageId") != id: | |
segment = data.get("segment", "") | |
content_buffer.write(segment) | |
else: | |
segment = data.get("segment", "") | |
think_buffer.write(segment) | |
think_state = 0 | |
else: | |
segment = data.get("segment", "") | |
content_buffer.write(segment) | |
except Exception as e: | |
print(f"处理响应出错: {e}") | |
think_content = think_buffer.getvalue() | |
response_content = content_buffer.getvalue() | |
# ---- Token 计算移到这里 ---- | |
print("\n----- 计算 Tokens (非流式, think模式) -----") | |
prompt_tokens, calculation_method = num_tokens_from_string(message, model) | |
print(f"输入 token 数: {prompt_tokens} (方法: {calculation_method})") | |
# 计算输出token并更新统计信息 | |
completion_tokens, comp_calc_method = num_tokens_from_string(think_content + response_content, model) | |
print(f"输出 token 数: {completion_tokens} (方法: {comp_calc_method})") | |
# 决定最终使用的计算方法 | |
final_calculation_method = "精确" if calculation_method == "精确" and comp_calc_method == "精确" else "估算" | |
# ---- Token 计算结束 ---- | |
# 保存对话历史并获取计算点数 | |
_, compute_points_used = save_conversation_history(session, cookies, session_token, conversation_id) | |
# 更新统计信息 | |
# 注意: 如果客户端在请求完成前断开连接(理论上非流式不太可能,但网络异常可能发生), | |
# Token 计算的准确性取决于 response 是否完整接收。Abacus 点数扣除不受影响。 | |
update_model_stats(model, prompt_tokens, completion_tokens, final_calculation_method, compute_points_used) | |
return jsonify({ | |
"id": f"chatcmpl-{str(uuid.uuid4())}", | |
"object": "chat.completion", | |
"created": int(time.time()), | |
"model": model, | |
"choices": [{ | |
"index": 0, | |
"message": { | |
"role": "assistant", | |
"content": f"<think>\n{think_content}\n</think>\n{response_content}" | |
}, | |
"finish_reason": "stop" | |
}], | |
"usage": { | |
"prompt_tokens": prompt_tokens, | |
"completion_tokens": completion_tokens, | |
"total_tokens": prompt_tokens + completion_tokens | |
} | |
}) | |
else: | |
for line in response.iter_lines(): | |
if line: | |
decoded_line = line.decode("utf-8") | |
segment = extract_segment(decoded_line) | |
if segment: | |
buffer.write(segment) | |
response_content = buffer.getvalue() | |
# ---- Token 计算移到这里 ---- | |
print("\n----- 计算 Tokens (非流式) -----") | |
prompt_tokens, calculation_method = num_tokens_from_string(message, model) | |
print(f"输入 token 数: {prompt_tokens} (方法: {calculation_method})") | |
# 计算输出token并更新统计信息 | |
completion_tokens, comp_calc_method = num_tokens_from_string(response_content, model) | |
print(f"输出 token 数: {completion_tokens} (方法: {comp_calc_method})") | |
# 决定最终使用的计算方法 | |
final_calculation_method = "精确" if calculation_method == "精确" and comp_calc_method == "精确" else "估算" | |
# ---- Token 计算结束 ---- | |
# 保存对话历史并获取计算点数 | |
_, compute_points_used = save_conversation_history(session, cookies, session_token, conversation_id) | |
# 更新统计信息 | |
# 注意: 如果客户端在请求完成前断开连接(理论上非流式不太可能,但网络异常可能发生), | |
# Token 计算的准确性取决于 response 是否完整接收。Abacus 点数扣除不受影响。 | |
update_model_stats(model, prompt_tokens, completion_tokens, final_calculation_method, compute_points_used) | |
return jsonify({ | |
"id": f"chatcmpl-{str(uuid.uuid4())}", | |
"object": "chat.completion", | |
"created": int(time.time()), | |
"model": model, | |
"choices": [{ | |
"index": 0, | |
"message": { | |
"role": "assistant", | |
"content": response_content | |
}, | |
"finish_reason": "stop" | |
}], | |
"usage": { | |
"prompt_tokens": prompt_tokens, | |
"completion_tokens": completion_tokens, | |
"total_tokens": prompt_tokens + completion_tokens | |
} | |
}) | |
except requests.exceptions.RequestException as e: | |
error_details = str(e) | |
if hasattr(e, 'response') and e.response is not None: | |
if hasattr(e.response, 'text'): | |
error_details += f" - Response: {e.response.text[:200]}" | |
print(f"发送消息失败: {error_details}") | |
# 如果是使用现有会话失败,尝试创建新会话重试一次 | |
if is_existing: | |
print("使用现有会话失败,尝试创建新会话...") | |
# 创建新会话 | |
deployment_id = "14b2a314cc" | |
external_app_id = model_map[model][0] if model in model_map and len(model_map[model]) >= 2 else None | |
if external_app_id: | |
new_conversation_id = create_conversation( | |
session, cookies, session_token, | |
external_application_id=external_app_id, | |
deployment_id=deployment_id | |
) | |
if new_conversation_id: | |
print(f"成功创建新会话ID: {new_conversation_id},重试请求") | |
# 更新全局存储的会话ID | |
global USER_DATA, CURRENT_USER | |
session, cookies, session_token, _, model_map, _ = USER_DATA[CURRENT_USER] | |
USER_DATA[CURRENT_USER] = (session, cookies, session_token, new_conversation_id, model_map, user_index) | |
# 保存到配置文件 | |
update_conversation_id(user_index, new_conversation_id) | |
# 修改payload使用新会话ID,并移除regenerate和editPrompt | |
payload["deploymentConversationId"] = new_conversation_id | |
if "regenerate" in payload: | |
del payload["regenerate"] | |
if "editPrompt" in payload: | |
del payload["editPrompt"] | |
try: | |
# 非流式重试逻辑与流式类似,但需要重新提取响应内容 | |
response = session.post( | |
CHAT_URL, | |
headers=headers, | |
data=json.dumps(payload), | |
stream=True, | |
cookies=None | |
) | |
response.raise_for_status() | |
# 重用现有提取逻辑... | |
# 但这里代码重复太多,实际应该重构为共享函数 | |
buffer = io.StringIO() | |
for line in response.iter_lines(): | |
if line: | |
decoded_line = line.decode("utf-8") | |
segment = extract_segment(decoded_line) | |
if segment: | |
buffer.write(segment) | |
response_content = buffer.getvalue() | |
# ---- 重试逻辑中的 Token 计算 ---- | |
print("\n----- 计算 Tokens (重试成功后) -----") | |
prompt_tokens, calculation_method = num_tokens_from_string(message, model) | |
print(f"输入 token 数: {prompt_tokens} (方法: {calculation_method})") | |
# 计算输出token并更新统计信息 | |
completion_tokens, comp_calc_method = num_tokens_from_string(response_content, model) | |
print(f"输出 token 数: {completion_tokens} (方法: {comp_calc_method})") | |
# 决定最终使用的计算方法 | |
final_calculation_method = "精确" if calculation_method == "精确" and comp_calc_method == "精确" else "估算" | |
# ---- Token 计算结束 ---- | |
# 保存对话历史并获取计算点数 | |
_, compute_points_used = save_conversation_history(session, cookies, session_token, new_conversation_id) | |
# 更新统计信息 | |
# 注意: 重试逻辑。Token 计算准确性依赖于 response 完整性。 | |
update_model_stats(model, prompt_tokens, completion_tokens, final_calculation_method, compute_points_used) | |
return jsonify({ | |
"id": f"chatcmpl-{str(uuid.uuid4())}", | |
"object": "chat.completion", | |
"created": int(time.time()), | |
"model": model, | |
"choices": [{ | |
"index": 0, | |
"message": { | |
"role": "assistant", | |
"content": response_content | |
}, | |
"finish_reason": "stop" | |
}], | |
"usage": { | |
"prompt_tokens": prompt_tokens, | |
"completion_tokens": completion_tokens, | |
"total_tokens": prompt_tokens + completion_tokens | |
} | |
}) | |
except Exception as retry_e: | |
print(f"重试失败: {retry_e}") | |
return jsonify({"error": f"Failed to send message: {error_details}"}), 500 | |
def format_message(messages): | |
buffer = io.StringIO() | |
role_map, prefix, messages = extract_role(messages) | |
for message in messages: | |
role = message.get("role") | |
role = "\b" + role_map[role] if prefix else role_map[role] | |
content = message.get("content").replace("\\n", "\n") | |
pattern = re.compile(r"<\|removeRole\|>\n") | |
if pattern.match(content): | |
content = pattern.sub("", content) | |
buffer.write(f"{content}\n") | |
else: | |
buffer.write(f"{role}: {content}\n\n") | |
formatted_message = buffer.getvalue() | |
return formatted_message | |
def extract_role(messages): | |
role_map = {"user": "Human", "assistant": "Assistant", "system": "System"} | |
prefix = True # 默认添加前缀 | |
first_message = messages[0]["content"] | |
pattern = re.compile( | |
r""" | |
<roleInfo>\s* | |
(?:user:\s*(?P<user>[^\n]*)\s*)? # Make user optional | |
(?:assistant:\s*(?P<assistant>[^\n]*)\s*)? # Make assistant optional | |
(?:system:\s*(?P<system>[^\n]*)\s*)? # Make system optional | |
(?:prefix:\s*(?P<prefix>[^\n]*)\s*)? # Make prefix optional | |
</roleInfo>\n | |
""", | |
re.VERBOSE, | |
) | |
match = pattern.search(first_message) | |
if match: | |
# 更新 role_map 如果提供了值 | |
user_role = match.group("user") | |
assistant_role = match.group("assistant") | |
system_role = match.group("system") | |
if user_role: role_map["user"] = user_role | |
if assistant_role: role_map["assistant"] = assistant_role | |
if system_role: role_map["system"] = system_role | |
# 检查 prefix 值:仅当显式设置为非 "1" 时才将 prefix 设为 False | |
prefix_value = match.group("prefix") | |
if prefix_value is not None and prefix_value != "1": | |
prefix = False | |
# 如果 prefix_value 是 None (标签不存在) 或 "1", prefix 保持 True | |
messages[0]["content"] = pattern.sub("", first_message) | |
print(f"Extracted role map:") | |
print( | |
f"User: {role_map['user']}, Assistant: {role_map['assistant']}, System: {role_map['system']}" | |
) | |
print(f"Using prefix: {prefix}") # 打印语句保持不变,反映最终结果 | |
# 如果没有匹配到 <roleInfo>,prefix 保持默认值 True | |
return (role_map, prefix, messages) | |
def health_check(): | |
global health_check_counter | |
health_check_counter += 1 | |
return jsonify({ | |
"status": "healthy", | |
"timestamp": datetime.now().isoformat(), | |
"checks": health_check_counter | |
}) | |
def keep_alive(): | |
"""每20分钟进行一次自我健康检查""" | |
while True: | |
try: | |
requests.get("http://127.0.0.1:7860/health") | |
time.sleep(1200) # 20分钟 | |
except: | |
pass # 忽略错误,保持运行 | |
def index(): | |
# 如果需要密码且用户未登录,重定向到登录页面 | |
if PASSWORD and not flask_session.get('logged_in'): | |
return redirect(url_for('login')) | |
# 否则重定向到仪表盘 | |
return redirect(url_for('dashboard')) | |
def num_tokens_from_string(string, model=""): | |
try: | |
print("\n===================== 开始计算token =====================") | |
print(f"模型: {model}") | |
print(f"输入内容长度: {len(string)} 字符") | |
request_data = { | |
"model": model, | |
"messages": [{"role": "user", "content": string}] | |
} | |
print(f"发送请求到tokenizer服务: {TOKENIZER_SERVICE_URL}") | |
print(f"请求数据: {json.dumps(request_data, ensure_ascii=False)}") | |
response = requests.post( | |
TOKENIZER_SERVICE_URL, | |
json=request_data, | |
timeout=10 | |
) | |
print(f"\nTokenizer响应状态码: {response.status_code}") | |
print(f"Tokenizer响应内容: {response.text}") | |
if response.status_code == 200: | |
result = response.json() | |
input_tokens = result.get("input_tokens", 0) | |
print(f"\n成功获取token数: {input_tokens}") | |
print(f"使用计算方法: 精确") | |
print("===================== 计算完成 =====================\n") | |
return input_tokens, "精确" | |
else: | |
estimated_tokens = len(string) // 4 | |
print(f"\nTokenizer服务错误: {response.status_code}") | |
print(f"错误响应: {response.text}") | |
print(f"使用估算token数: {estimated_tokens}") | |
print(f"使用计算方法: 估算") | |
print("===================== 计算完成 =====================\n") | |
return estimated_tokens, "估算" | |
except Exception as e: | |
estimated_tokens = len(string) // 4 | |
print(f"\n计算token时发生错误: {str(e)}") | |
print(f"使用估算token数: {estimated_tokens}") | |
print(f"使用计算方法: 估算") | |
print("===================== 计算完成 =====================\n") | |
return estimated_tokens, "估算" | |
# 更新模型使用统计 | |
def update_model_stats(model, prompt_tokens, completion_tokens, calculation_method="estimate", compute_points=None): | |
global model_usage_stats, total_tokens, model_usage_records | |
# 添加调用记录 | |
# 获取UTC时间 | |
utc_now = datetime.utcnow() | |
# 转换为北京时间 (UTC+8) | |
beijing_time = utc_now + timedelta(hours=8) | |
call_time = beijing_time.strftime('%Y-%m-%d %H:%M:%S') # 北京时间 | |
record = { | |
"model": model, | |
"call_time": call_time, | |
"prompt_tokens": prompt_tokens, | |
"completion_tokens": completion_tokens, | |
"calculation_method": calculation_method, # 直接使用传入的值 | |
"compute_points": compute_points | |
} | |
model_usage_records.append(record) | |
# 限制记录数量,保留最新的500条 | |
if len(model_usage_records) > 500: | |
model_usage_records.pop(0) | |
# 保存调用记录到本地文件 | |
save_model_usage_records() | |
# 更新聚合统计 | |
if model not in model_usage_stats: | |
model_usage_stats[model] = { | |
"count": 0, | |
"prompt_tokens": 0, | |
"completion_tokens": 0, | |
"total_tokens": 0 | |
} | |
model_usage_stats[model]["count"] += 1 | |
model_usage_stats[model]["prompt_tokens"] += prompt_tokens | |
model_usage_stats[model]["completion_tokens"] += completion_tokens | |
model_usage_stats[model]["total_tokens"] += (prompt_tokens + completion_tokens) | |
total_tokens["prompt"] += prompt_tokens | |
total_tokens["completion"] += completion_tokens | |
total_tokens["total"] += (prompt_tokens + completion_tokens) | |
# 获取计算点信息 | |
def get_compute_points(): | |
global compute_points, USER_DATA, users_compute_points | |
if USER_NUM == 0: | |
return | |
# 清空用户计算点列表 | |
users_compute_points = [] | |
# 累计总计算点 | |
total_left = 0 | |
total_points = 0 | |
# 获取每个用户的计算点信息 | |
for i, user_data in enumerate(USER_DATA): | |
try: | |
session, cookies, session_token, _, _, _ = user_data | |
# 检查token是否有效 | |
if is_token_expired(session_token): | |
session_token = refresh_token(session, cookies) | |
if not session_token: | |
print(f"用户{i+1}刷新token失败,无法获取计算点信息") | |
continue | |
USER_DATA[i] = (session, cookies, session_token, user_data[3], user_data[4], i) | |
headers = { | |
"accept": "application/json, text/plain, */*", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"baggage": f"sentry-environment=production,sentry-release=93da8385541a6ce339b1f41b0c94428c70657e22,sentry-public_key=3476ea6df1585dd10e92cdae3a66ff49,sentry-trace_id={TRACE_ID}", | |
"reai-ui": "1", | |
"sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"", | |
"sec-ch-ua-mobile": "?0", | |
"sec-ch-ua-platform": "\"Windows\"", | |
"sec-fetch-dest": "empty", | |
"sec-fetch-mode": "cors", | |
"sec-fetch-site": "same-origin", | |
"sentry-trace": SENTRY_TRACE, | |
"session-token": session_token, | |
"x-abacus-org-host": "apps", | |
"cookie": cookies | |
} | |
response = session.get( | |
COMPUTE_POINTS_URL, | |
headers=headers | |
) | |
if response.status_code == 200: | |
result = response.json() | |
if result.get("success") and "result" in result: | |
data = result["result"] | |
left = data.get("computePointsLeft", 0) | |
total = data.get("totalComputePoints", 0) | |
used = total - left | |
percentage = round((used / total) * 100, 2) if total > 0 else 0 | |
# 获取北京时间 | |
beijing_now = datetime.utcnow() + timedelta(hours=8) | |
# 添加到用户列表 | |
user_points = { | |
"user_id": i + 1, # 用户ID从1开始 | |
"left": left, | |
"total": total, | |
"used": used, | |
"percentage": percentage, | |
"last_update": beijing_now | |
} | |
users_compute_points.append(user_points) | |
# 累计总数 | |
total_left += left | |
total_points += total | |
print(f"用户{i+1}计算点信息更新成功: 剩余 {left}, 总计 {total}") | |
# 对于第一个用户,获取计算点使用日志 | |
if i == 0: | |
get_compute_points_log(session, cookies, session_token) | |
else: | |
print(f"获取用户{i+1}计算点信息失败: {result.get('error', '未知错误')}") | |
else: | |
print(f"获取用户{i+1}计算点信息失败,状态码: {response.status_code}") | |
except Exception as e: | |
print(f"获取用户{i+1}计算点信息异常: {e}") | |
# 更新全局计算点信息(所有用户总和) | |
if users_compute_points: | |
compute_points["left"] = total_left | |
compute_points["total"] = total_points | |
compute_points["used"] = total_points - total_left | |
compute_points["percentage"] = round((compute_points["used"] / compute_points["total"]) * 100, 2) if compute_points["total"] > 0 else 0 | |
compute_points["last_update"] = datetime.utcnow() + timedelta(hours=8) # 北京时间 | |
print(f"所有用户计算点总计: 剩余 {total_left}, 总计 {total_points}") | |
# 获取计算点使用日志 | |
def get_compute_points_log(session, cookies, session_token): | |
global compute_points_log | |
try: | |
headers = { | |
"accept": "application/json, text/plain, */*", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"content-type": "application/json", | |
"reai-ui": "1", | |
"sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"", | |
"sec-ch-ua-mobile": "?0", | |
"sec-ch-ua-platform": "\"Windows\"", | |
"sec-fetch-dest": "empty", | |
"sec-fetch-mode": "cors", | |
"sec-fetch-site": "same-site", | |
"session-token": session_token, | |
"x-abacus-org-host": "apps", | |
"cookie": cookies | |
} | |
response = session.post( | |
COMPUTE_POINTS_LOG_URL, | |
headers=headers, | |
json={"byLlm": True} | |
) | |
if response.status_code == 200: | |
result = response.json() | |
if result.get("success") and "result" in result: | |
data = result["result"] | |
compute_points_log["columns"] = data.get("columns", {}) | |
compute_points_log["log"] = data.get("log", []) | |
print(f"计算点使用日志更新成功,获取到 {len(compute_points_log['log'])} 条记录") | |
else: | |
print(f"获取计算点使用日志失败: {result.get('error', '未知错误')}") | |
else: | |
print(f"获取计算点使用日志失败,状态码: {response.status_code}") | |
except Exception as e: | |
print(f"获取计算点使用日志异常: {e}") | |
# 添加登录相关路由 | |
def login(): | |
error = None | |
if request.method == "POST": | |
password = request.form.get("password") | |
if password and hashlib.sha256(password.encode()).hexdigest() == PASSWORD: | |
flask_session['logged_in'] = True | |
flask_session.permanent = True | |
return redirect(url_for('dashboard')) | |
else: | |
# 密码错误时提示使用环境变量密码 | |
error = "密码不正确。请使用设置的环境变量 password 或 password.txt 中的值作为密码和API认证密钥。" | |
# 传递空间URL给模板 | |
return render_template('login.html', error=error, space_url=SPACE_URL) | |
def logout(): | |
flask_session.clear() | |
return redirect(url_for('login')) | |
def dashboard(): | |
# 在每次访问仪表盘时更新计算点信息 | |
get_compute_points() | |
# 计算运行时间(使用北京时间) | |
beijing_now = datetime.utcnow() + timedelta(hours=8) | |
uptime = beijing_now - START_TIME | |
days = uptime.days | |
hours, remainder = divmod(uptime.seconds, 3600) | |
minutes, seconds = divmod(remainder, 60) | |
if days > 0: | |
uptime_str = f"{days}天 {hours}小时 {minutes}分钟" | |
elif hours > 0: | |
uptime_str = f"{hours}小时 {minutes}分钟" | |
else: | |
uptime_str = f"{minutes}分钟 {seconds}秒" | |
# 当前北京年份 | |
beijing_year = beijing_now.year | |
return render_template( | |
'dashboard.html', | |
uptime=uptime_str, | |
health_checks=health_check_counter, | |
user_count=USER_NUM, | |
models=sorted(list(MODELS)), | |
year=beijing_year, | |
model_stats=model_usage_stats, | |
total_tokens=total_tokens, | |
compute_points=compute_points, | |
compute_points_log=compute_points_log, | |
space_url=SPACE_URL, # 传递空间URL | |
users_compute_points=users_compute_points, # 传递用户计算点信息 | |
model_usage_records=model_usage_records, # 传递模型使用记录 | |
) | |
# 添加更新计算点数记录设置的路由 | |
def update_compute_point_toggle(): | |
try: | |
(session, cookies, session_token, conversation_id, model_map, user_index) = get_user_data() | |
data = request.get_json() | |
if data and "always_display" in data: | |
headers = { | |
"accept": "application/json, text/plain, */*", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"content-type": "application/json", | |
"reai-ui": "1", | |
"sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"", | |
"sec-ch-ua-mobile": "?0", | |
"sec-ch-ua-platform": "\"Windows\"", | |
"sec-fetch-dest": "empty", | |
"sec-fetch-mode": "cors", | |
"sec-fetch-site": "same-site", | |
"x-abacus-org-host": "apps" | |
} | |
if session_token: | |
headers["session-token"] = session_token | |
response = session.post( | |
COMPUTE_POINT_TOGGLE_URL, | |
headers=headers, | |
json={"alwaysDisplay": data["always_display"]}, | |
cookies=None | |
) | |
if response.status_code == 200: | |
result = response.json() | |
if result.get("success"): | |
print(f"更新计算点数记录设置为: {data['always_display']}") | |
return jsonify({"success": True}) | |
return jsonify({"success": False, "error": "API调用失败"}) | |
else: | |
return jsonify({"success": False, "error": "缺少always_display参数"}) | |
except Exception as e: | |
print(f"更新计算点数记录设置失败: {e}") | |
return jsonify({"success": False, "error": str(e)}) | |
# 获取Hugging Face Space URL | |
def get_space_url(): | |
# 尝试从环境变量获取 | |
space_url = os.environ.get("SPACE_URL") | |
if space_url: | |
return space_url | |
# 如果SPACE_URL不存在,尝试从SPACE_ID构建 | |
space_id = os.environ.get("SPACE_ID") | |
if space_id: | |
username, space_name = space_id.split("/") | |
# 将空间名称中的下划线替换为连字符 | |
# 注意:Hugging Face生成的URL会自动将空间名称中的下划线(_)替换为连字符(-) | |
# 例如:"abacus_chat_proxy" 会变成 "abacus-chat-proxy" | |
space_name = space_name.replace("_", "-") | |
return f"https://{username}-{space_name}.hf.space" | |
# 如果以上都不存在,尝试从单独的用户名和空间名构建 | |
username = os.environ.get("SPACE_USERNAME") | |
space_name = os.environ.get("SPACE_NAME") | |
if username and space_name: | |
# 将空间名称中的下划线替换为连字符 | |
# 同上,Hugging Face会自动进行此转换 | |
space_name = space_name.replace("_", "-") | |
return f"https://{username}-{space_name}.hf.space" | |
# 默认返回None | |
return None | |
# 获取空间URL | |
SPACE_URL = get_space_url() | |
if SPACE_URL: | |
print(f"Space URL: {SPACE_URL}") | |
print("注意:Hugging Face生成的URL会自动将空间名称中的下划线(_)替换为连字符(-)") | |
def save_conversation_history(session, cookies, session_token, conversation_id, deployment_id="14b2a314cc"): | |
"""保存对话历史,返回使用的计算点数""" | |
if not conversation_id: | |
return False, None | |
headers = { | |
"accept": "application/json, text/plain, */*", | |
"accept-language": "zh-CN,zh;q=0.9", | |
"baggage": f"sentry-environment=production,sentry-release=946244517de08b08598b94f18098411f5a5352d5,sentry-public_key=3476ea6df1585dd10e92cdae3a66ff49,sentry-trace_id={TRACE_ID}", | |
"reai-ui": "1", | |
"sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"", | |
"sec-ch-ua-mobile": "?0", | |
"sec-ch-ua-platform": "\"Windows\"", | |
"sec-fetch-dest": "empty", | |
"sec-fetch-mode": "cors", | |
"sec-fetch-site": "same-origin", | |
"sentry-trace": f"{TRACE_ID}-800cb7f4613dec52", | |
"x-abacus-org-host": "apps" | |
} | |
if session_token: | |
headers["session-token"] = session_token | |
params = { | |
"deploymentId": deployment_id, | |
"deploymentConversationId": conversation_id, | |
"skipDocumentBoundingBoxes": "true", | |
"filterIntermediateConversationEvents": "false", | |
"getUnusedDocumentUploads": "false" | |
} | |
try: | |
response = session.get( | |
GET_CONVERSATION_URL, | |
headers=headers, | |
params=params, | |
cookies=None | |
) | |
if response.status_code == 200: | |
data = response.json() | |
if data.get("success"): | |
# 从最后一条BOT消息中获取计算点数 | |
history = data.get("result", {}).get("history", []) | |
compute_points = None | |
for msg in reversed(history): | |
if msg.get("role") == "BOT": | |
compute_points = msg.get("computePointsUsed") | |
break | |
print(f"成功保存对话历史: {conversation_id}, 使用计算点: {compute_points}") | |
return True, compute_points | |
else: | |
print(f"保存对话历史失败: {data.get('error', '未知错误')}") | |
else: | |
print(f"保存对话历史失败,状态码: {response.status_code}") | |
return False, None | |
except Exception as e: | |
print(f"保存对话历史时出错: {e}") | |
return False, None | |
if __name__ == "__main__": | |
# 启动保活线程 | |
threading.Thread(target=keep_alive, daemon=True).start() | |
# 加载历史模型调用记录 | |
load_model_usage_records() | |
# 获取初始计算点信息 | |
get_compute_points() | |
port = int(os.environ.get("PORT", 9876)) | |
app.run(port=port, host="0.0.0.0") | |