Ha123456/abacus_chat_proxy5
0
1from flask import Flask, request, jsonify, Response, render_template_string, render_template, redirect, url_for, session as flask_session2import requests3import time4import json5import uuid6import random7import io8import re9from functools import wraps10import hashlib11import jwt 12import os13import threading14from datetime import datetime, timedelta15import tiktoken # 导入tiktoken来计算token数量16 17app = Flask(__name__, template_folder='templates')18app.secret_key = os.environ.get("SECRET_KEY", "abacus_chat_proxy_secret_key")19app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(days=7)20 21 22API_ENDPOINT_URL = "https://abacus.ai/api/v0/describeDeployment"23MODEL_LIST_URL = "https://abacus.ai/api/v0/listExternalApplications"24CHAT_URL = "https://apps.abacus.ai/api/_chatLLMSendMessageSSE"25USER_INFO_URL = "https://abacus.ai/api/v0/_getUserInfo"26COMPUTE_POINTS_URL = "https://apps.abacus.ai/api/_getOrganizationComputePoints"27COMPUTE_POINTS_LOG_URL = "https://abacus.ai/api/v0/_getOrganizationComputePointLog"28 29 30USER_AGENTS = [31 "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/116.0.0.0 Safari/537.36"32]33 34 35PASSWORD = None36USER_NUM = 037USER_DATA = []38CURRENT_USER = -139MODELS = set()40 41 42TRACE_ID = "3042e28b3abf475d8d973c7e904935af"43SENTRY_TRACE = f"{TRACE_ID}-80d9d2538b2682d0"44 45 46# 添加一个计数器记录健康检查次数47health_check_counter = 048 49 50# 添加统计变量51model_usage_stats = {} # 模型使用次数统计52total_tokens = {53 "prompt": 0, # 输入token统计54 "completion": 0, # 输出token统计55 "total": 0 # 总token统计56}57 58# 模型调用记录59model_usage_records = [] # 每次调用详细记录60MODEL_USAGE_RECORDS_FILE = "model_usage_records.json" # 调用记录保存文件61 62# 计算点信息63compute_points = {64 "left": 0, # 剩余计算点65 "total": 0, # 总计算点66 "used": 0, # 已使用计算点67 "percentage": 0, # 使用百分比68 "last_update": None # 最后更新时间69}70 71# 计算点使用日志72compute_points_log = {73 "columns": {}, # 列名74 "log": [] # 日志数据75}76 77# 多用户计算点信息78users_compute_points = []79 80# 记录启动时间81START_TIME = datetime.utcnow() + timedelta(hours=8) # 北京时间82 83 84# 自定义JSON编码器，处理datetime对象85class DateTimeEncoder(json.JSONEncoder):86 def default(self, obj):87 if isinstance(obj, datetime):88 return obj.strftime('%Y-%m-%d %H:%M:%S')89 return super(DateTimeEncoder, self).default(obj)90 91 92# 加载模型调用记录93def load_model_usage_records():94 global model_usage_records95 try:96 if os.path.exists(MODEL_USAGE_RECORDS_FILE):97 with open(MODEL_USAGE_RECORDS_FILE, 'r', encoding='utf-8') as f:98 records = json.load(f)99 if isinstance(records, list):100 model_usage_records = records101 print(f"成功加载 {len(model_usage_records)} 条模型调用记录")102 else:103 print("调用记录文件格式不正确，初始化为空列表")104 except Exception as e:105 print(f"加载模型调用记录失败: {e}")106 model_usage_records = []107 108# 保存模型调用记录109def save_model_usage_records():110 try:111 with open(MODEL_USAGE_RECORDS_FILE, 'w', encoding='utf-8') as f:112 json.dump(model_usage_records, f, ensure_ascii=False, indent=2, cls=DateTimeEncoder)113 print(f"成功保存 {len(model_usage_records)} 条模型调用记录")114 except Exception as e:115 print(f"保存模型调用记录失败: {e}")116 117 118def resolve_config():119 # 从环境变量读取多组配置120 config_list = []121 i = 1122 while True:123 covid = os.environ.get(f"covid_{i}")124 cookie = os.environ.get(f"cookie_{i}")125 if not (covid and cookie):126 break127 config_list.append({128 "conversation_id": covid,129 "cookies": cookie130 })131 i += 1132 133 # 如果环境变量存在配置，使用环境变量的配置134 if config_list:135 return config_list136 137 # 如果环境变量不存在，从文件读取138 try:139 with open("config.json", "r") as f:140 config = json.load(f)141 config_list = config.get("config")142 return config_list143 except FileNotFoundError:144 print("未找到config.json文件")145 return []146 except json.JSONDecodeError:147 print("config.json格式错误")148 return []149 150 151def get_password():152 global PASSWORD153 # 从环境变量读取密码154 env_password = os.environ.get("password")155 if env_password:156 PASSWORD = hashlib.sha256(env_password.encode()).hexdigest()157 return158 159 # 如果环境变量不存在，从文件读取160 try:161 with open("password.txt", "r") as f:162 PASSWORD = f.read().strip()163 except FileNotFoundError:164 with open("password.txt", "w") as f:165 PASSWORD = None166 167 168def require_auth(f):169 @wraps(f)170 def decorated(*args, **kwargs):171 if not PASSWORD:172 return f(*args, **kwargs)173 174 # 检查Flask会话是否已登录175 if flask_session.get('logged_in'):176 return f(*args, **kwargs)177 178 # 如果是API请求，检查Authorization头179 auth = request.authorization180 if not auth or not check_auth(auth.token):181 # 如果是浏览器请求，重定向到登录页面182 if request.headers.get('Accept', '').find('text/html') >= 0:183 return redirect(url_for('login'))184 return jsonify({"error": "Unauthorized access"}), 401185 return f(*args, **kwargs)186 187 return decorated188 189 190def check_auth(token):191 return hashlib.sha256(token.encode()).hexdigest() == PASSWORD192 193 194def is_token_expired(token):195 if not token:196 return True197 198 try:199 # Malkodi tokenon sen validigo de subskribo200 payload = jwt.decode(token, options={"verify_signature": False})201 # Akiru eksvalidiĝan tempon, konsideru eksvalidiĝinta 5 minutojn antaŭe202 return payload.get('exp', 0) - time.time() < 300203 except:204 return True205 206 207def refresh_token(session, cookies):208 """Uzu kuketon por refreŝigi session token, nur revenigu novan tokenon"""209 headers = {210 "accept": "application/json, text/plain, */*",211 "accept-language": "zh-CN,zh;q=0.9",212 "content-type": "application/json",213 "reai-ui": "1",214 "sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"",215 "sec-ch-ua-mobile": "?0",216 "sec-ch-ua-platform": "\"Windows\"",217 "sec-fetch-dest": "empty",218 "sec-fetch-mode": "cors",219 "sec-fetch-site": "same-site",220 "x-abacus-org-host": "apps",221 "user-agent": random.choice(USER_AGENTS),222 "origin": "https://apps.abacus.ai",223 "referer": "https://apps.abacus.ai/",224 "cookie": cookies225 }226 227 try:228 response = session.post(229 USER_INFO_URL,230 headers=headers,231 json={},232 cookies=None233 )234 235 if response.status_code == 200:236 response_data = response.json()237 if response_data.get('success') and 'sessionToken' in response_data.get('result', {}):238 return response_data['result']['sessionToken']239 else:240 print(f"刷新token失败: {response_data.get('error', '未知错误')}")241 return None242 else:243 print(f"刷新token失败，状态码: {response.status_code}")244 return None245 except Exception as e:246 print(f"刷新token异常: {e}")247 return None248 249 250def get_model_map(session, cookies, session_token):251 """Akiru disponeblan modelan liston kaj ĝiajn mapajn rilatojn"""252 headers = {253 "accept": "application/json, text/plain, */*",254 "accept-language": "zh-CN,zh;q=0.9",255 "content-type": "application/json",256 "reai-ui": "1",257 "sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"",258 "sec-ch-ua-mobile": "?0",259 "sec-ch-ua-platform": "\"Windows\"",260 "sec-fetch-dest": "empty",261 "sec-fetch-mode": "cors",262 "sec-fetch-site": "same-site",263 "x-abacus-org-host": "apps",264 "user-agent": random.choice(USER_AGENTS),265 "origin": "https://apps.abacus.ai",266 "referer": "https://apps.abacus.ai/",267 "cookie": cookies268 }269 270 if session_token:271 headers["session-token"] = session_token272 273 model_map = {}274 models_set = set()275 276 try:277 response = session.post(278 MODEL_LIST_URL,279 headers=headers,280 json={},281 cookies=None282 )283 284 if response.status_code != 200:285 print(f"获取模型列表失败，状态码: {response.status_code}")286 raise Exception("API请求失败")287 288 data = response.json()289 if not data.get('success'):290 print(f"获取模型列表失败: {data.get('error', '未知错误')}")291 raise Exception("API返回错误")292 293 applications = []294 if isinstance(data.get('result'), dict):295 applications = data.get('result', {}).get('externalApplications', [])296 elif isinstance(data.get('result'), list):297 applications = data.get('result', [])298 299 for app in applications:300 app_name = app.get('name', '')301 app_id = app.get('externalApplicationId', '')302 prediction_overrides = app.get('predictionOverrides', {})303 llm_name = prediction_overrides.get('llmName', '') if prediction_overrides else ''304 305 if not (app_name and app_id and llm_name):306 continue307 308 model_name = app_name309 model_map[model_name] = (app_id, llm_name)310 models_set.add(model_name)311 312 if not model_map:313 raise Exception("未找到任何可用模型")314 315 return model_map, models_set316 317 except Exception as e:318 print(f"获取模型列表异常: {e}")319 raise320 321 322def init_session():323 get_password()324 global USER_NUM, MODELS, USER_DATA325 config_list = resolve_config()326 user_num = len(config_list)327 all_models = set()328 329 for i in range(user_num):330 user = config_list[i]331 cookies = user.get("cookies")332 conversation_id = user.get("conversation_id")333 session = requests.Session()334 335 session_token = refresh_token(session, cookies)336 if not session_token:337 print(f"无法获取cookie {i+1}的token")338 continue339 340 try:341 model_map, models_set = get_model_map(session, cookies, session_token)342 all_models.update(models_set)343 USER_DATA.append((session, cookies, session_token, conversation_id, model_map))344 except Exception as e:345 print(f"配置用户 {i+1} 失败: {e}")346 continue347 348 USER_NUM = len(USER_DATA)349 if USER_NUM == 0:350 print("No user available, exiting...")351 exit(1)352 353 MODELS = all_models354 print(f"启动完成，共配置 {USER_NUM} 个用户")355 356 357def update_cookie(session, cookies):358 cookie_jar = {}359 for key, value in session.cookies.items():360 cookie_jar[key] = value361 cookie_dict = {}362 for item in cookies.split(";"):363 key, value = item.strip().split("=", 1)364 cookie_dict[key] = value365 cookie_dict.update(cookie_jar)366 cookies = "; ".join([f"{key}={value}" for key, value in cookie_dict.items()])367 return cookies368 369 370user_data = init_session()371 372 373@app.route("/v1/models", methods=["GET"])374@require_auth375def get_models():376 if len(MODELS) == 0:377 return jsonify({"error": "No models available"}), 500378 model_list = []379 for model in MODELS:380 model_list.append(381 {382 "id": model,383 "object": "model",384 "created": int(time.time()),385 "owned_by": "Elbert",386 "name": model,387 }388 )389 return jsonify({"object": "list", "data": model_list})390 391 392@app.route("/v1/chat/completions", methods=["POST"])393@require_auth394def chat_completions():395 openai_request = request.get_json()396 stream = openai_request.get("stream", False)397 messages = openai_request.get("messages")398 if messages is None:399 return jsonify({"error": "Messages is required", "status": 400}), 400400 model = openai_request.get("model")401 if model not in MODELS:402 return (403 jsonify(404 {405 "error": "Model not available, check if it is configured properly",406 "status": 404,407 }408 ),409 404,410 )411 message = format_message(messages)412 think = (413 openai_request.get("think", False) if model == "Claude Sonnet 3.7" else False414 )415 return (416 send_message(message, model, think)417 if stream418 else send_message_non_stream(message, model, think)419 )420 421 422def get_user_data():423 global CURRENT_USER424 CURRENT_USER = (CURRENT_USER + 1) % USER_NUM425 print(f"使用配置 {CURRENT_USER+1}")426 427 # Akiru uzantajn datumojn428 session, cookies, session_token, conversation_id, model_map = USER_DATA[CURRENT_USER]429 430 # Kontrolu ĉu la tokeno eksvalidiĝis, se jes, refreŝigu ĝin431 if is_token_expired(session_token):432 print(f"Cookie {CURRENT_USER+1}的token已过期或即将过期，正在刷新...")433 new_token = refresh_token(session, cookies)434 if new_token:435 # Ĝisdatigu la globale konservitan tokenon436 USER_DATA[CURRENT_USER] = (session, cookies, new_token, conversation_id, model_map)437 session_token = new_token438 print(f"成功更新token: {session_token[:15]}...{session_token[-15:]}")439 else:440 print(f"警告：无法刷新Cookie {CURRENT_USER+1}的token，继续使用当前token")441 442 return (session, cookies, session_token, conversation_id, model_map)443 444 445def generate_trace_id():446 """Generu novan trace_id kaj sentry_trace"""447 trace_id = str(uuid.uuid4()).replace('-', '')448 sentry_trace = f"{trace_id}-{str(uuid.uuid4())[:16]}"449 return trace_id, sentry_trace450 451 452def send_message(message, model, think=False):453 """Flua traktado kaj plusendo de mesaĝoj"""454 (session, cookies, session_token, conversation_id, model_map) = get_user_data()455 trace_id, sentry_trace = generate_trace_id()456 457 # 计算输入token458 prompt_tokens = num_tokens_from_string(message)459 completion_buffer = io.StringIO() # 收集所有输出用于计算token460 461 headers = {462 "accept": "text/event-stream",463 "accept-language": "zh-CN,zh;q=0.9",464 "baggage": f"sentry-environment=production,sentry-release=975eec6685013679c139fc88db2c48e123d5c604,sentry-public_key=3476ea6df1585dd10e92cdae3a66ff49,sentry-trace_id={trace_id}",465 "content-type": "text/plain;charset=UTF-8",466 "cookie": cookies,467 "sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"",468 "sec-ch-ua-mobile": "?0",469 "sec-ch-ua-platform": "\"Windows\"",470 "sec-fetch-dest": "empty",471 "sec-fetch-mode": "cors",472 "sec-fetch-site": "same-origin",473 "sentry-trace": sentry_trace,474 "user-agent": random.choice(USER_AGENTS)475 }476 477 if session_token:478 headers["session-token"] = session_token479 480 payload = {481 "requestId": str(uuid.uuid4()),482 "deploymentConversationId": conversation_id,483 "message": message,484 "isDesktop": False,485 "chatConfig": {486 "timezone": "Asia/Shanghai",487 "language": "zh-CN"488 },489 "llmName": model_map[model][1],490 "externalApplicationId": model_map[model][0],491 "regenerate": True,492 "editPrompt": True493 }494 495 if think:496 payload["useThinking"] = think497 498 try:499 response = session.post(500 CHAT_URL,501 headers=headers,502 data=json.dumps(payload),503 stream=True504 )505 506 response.raise_for_status()507 508 def extract_segment(line_data):509 try:510 data = json.loads(line_data)511 if "segment" in data:512 if isinstance(data["segment"], str):513 return data["segment"]514 elif isinstance(data["segment"], dict) and "segment" in data["segment"]:515 return data["segment"]["segment"]516 return ""517 except:518 return ""519 520 def generate():521 id = ""522 think_state = 2523 524 yield "data: " + json.dumps({"object": "chat.completion.chunk", "choices": [{"delta": {"role": "assistant"}}]}) + "\n\n"525 526 for line in response.iter_lines():527 if line:528 decoded_line = line.decode("utf-8")529 try:530 if think:531 data = json.loads(decoded_line)532 if data.get("type") != "text":533 continue534 elif think_state == 2:535 id = data.get("messageId")536 segment = "<think>\n" + data.get("segment", "")537 completion_buffer.write(segment) # 收集输出538 yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n"539 think_state = 1540 elif think_state == 1:541 if data.get("messageId") != id:542 segment = data.get("segment", "")543 completion_buffer.write(segment) # 收集输出544 yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n"545 else:546 segment = "\n</think>\n" + data.get("segment", "")547 completion_buffer.write(segment) # 收集输出548 yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n"549 think_state = 0550 else:551 segment = data.get("segment", "")552 completion_buffer.write(segment) # 收集输出553 yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n"554 else:555 segment = extract_segment(decoded_line)556 if segment:557 completion_buffer.write(segment) # 收集输出558 yield f"data: {json.dumps({'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': segment}}]})}\n\n"559 except Exception as e:560 print(f"处理响应出错: {e}")561 562 yield "data: " + json.dumps({"object": "chat.completion.chunk", "choices": [{"delta": {}, "finish_reason": "stop"}]}) + "\n\n"563 yield "data: [DONE]\n\n"564 565 # 在流式传输完成后计算token并更新统计566 completion_tokens = num_tokens_from_string(completion_buffer.getvalue())567 update_model_stats(model, prompt_tokens, completion_tokens)568 569 return Response(generate(), mimetype="text/event-stream")570 except requests.exceptions.RequestException as e:571 error_details = str(e)572 if hasattr(e, 'response') and e.response is not None:573 if hasattr(e.response, 'text'):574 error_details += f" - Response: {e.response.text[:200]}"575 print(f"发送消息失败: {error_details}")576 return jsonify({"error": f"Failed to send message: {error_details}"}), 500577 578 579def send_message_non_stream(message, model, think=False):580 """Ne-flua traktado de mesaĝoj"""581 (session, cookies, session_token, conversation_id, model_map) = get_user_data()582 trace_id, sentry_trace = generate_trace_id()583 584 # 计算输入token585 prompt_tokens = num_tokens_from_string(message)586 587 headers = {588 "accept": "text/event-stream",589 "accept-language": "zh-CN,zh;q=0.9",590 "baggage": f"sentry-environment=production,sentry-release=975eec6685013679c139fc88db2c48e123d5c604,sentry-public_key=3476ea6df1585dd10e92cdae3a66ff49,sentry-trace_id={trace_id}",591 "content-type": "text/plain;charset=UTF-8",592 "cookie": cookies,593 "sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"",594 "sec-ch-ua-mobile": "?0",595 "sec-ch-ua-platform": "\"Windows\"",596 "sec-fetch-dest": "empty",597 "sec-fetch-mode": "cors",598 "sec-fetch-site": "same-origin",599 "sentry-trace": sentry_trace,600 "user-agent": random.choice(USER_AGENTS)601 }602 603 if session_token:604 headers["session-token"] = session_token605 606 payload = {607 "requestId": str(uuid.uuid4()),608 "deploymentConversationId": conversation_id,609 "message": message,610 "isDesktop": False,611 "chatConfig": {612 "timezone": "Asia/Shanghai",613 "language": "zh-CN"614 },615 "llmName": model_map[model][1],616 "externalApplicationId": model_map[model][0],617 "regenerate": True,618 "editPrompt": True619 }620 621 if think:622 payload["useThinking"] = think623 624 try:625 response = session.post(626 CHAT_URL,627 headers=headers,628 data=json.dumps(payload),629 stream=True630 )631 632 response.raise_for_status()633 buffer = io.StringIO()634 635 def extract_segment(line_data):636 try:637 data = json.loads(line_data)638 if "segment" in data:639 if isinstance(data["segment"], str):640 return data["segment"]641 elif isinstance(data["segment"], dict) and "segment" in data["segment"]:642 return data["segment"]["segment"]643 return ""644 except:645 return ""646 647 if think:648 id = ""649 think_state = 2650 think_buffer = io.StringIO()651 content_buffer = io.StringIO()652 653 for line in response.iter_lines():654 if line:655 decoded_line = line.decode("utf-8")656 try:657 data = json.loads(decoded_line)658 if data.get("type") != "text":659 continue660 elif think_state == 2:661 id = data.get("messageId")662 segment = data.get("segment", "")663 think_buffer.write(segment)664 think_state = 1665 elif think_state == 1:666 if data.get("messageId") != id:667 segment = data.get("segment", "")668 content_buffer.write(segment)669 else:670 segment = data.get("segment", "")671 think_buffer.write(segment)672 think_state = 0673 else:674 segment = data.get("segment", "")675 content_buffer.write(segment)676 except Exception as e:677 print(f"处理响应出错: {e}")678 679 think_content = think_buffer.getvalue()680 response_content = content_buffer.getvalue()681 682 # 计算输出token并更新统计信息683 completion_tokens = num_tokens_from_string(think_content + response_content)684 update_model_stats(model, prompt_tokens, completion_tokens)685 686 return jsonify({687 "id": f"chatcmpl-{str(uuid.uuid4())}",688 "object": "chat.completion",689 "created": int(time.time()),690 "model": model,691 "choices": [{692 "index": 0,693 "message": {694 "role": "assistant",695 "content": f"<think>\n{think_content}\n</think>\n{response_content}"696 },697 "finish_reason": "stop"698 }],699 "usage": {700 "prompt_tokens": prompt_tokens,701 "completion_tokens": completion_tokens,702 "total_tokens": prompt_tokens + completion_tokens703 }704 })705 else:706 for line in response.iter_lines():707 if line:708 decoded_line = line.decode("utf-8")709 segment = extract_segment(decoded_line)710 if segment:711 buffer.write(segment)712 713 response_content = buffer.getvalue()714 715 # 计算输出token并更新统计信息716 completion_tokens = num_tokens_from_string(response_content)717 update_model_stats(model, prompt_tokens, completion_tokens)718 719 return jsonify({720 "id": f"chatcmpl-{str(uuid.uuid4())}",721 "object": "chat.completion",722 "created": int(time.time()),723 "model": model,724 "choices": [{725 "index": 0,726 "message": {727 "role": "assistant",728 "content": response_content729 },730 "finish_reason": "stop"731 }],732 "usage": {733 "prompt_tokens": prompt_tokens,734 "completion_tokens": completion_tokens,735 "total_tokens": prompt_tokens + completion_tokens736 }737 })738 except requests.exceptions.RequestException as e:739 error_details = str(e)740 if hasattr(e, 'response') and e.response is not None:741 if hasattr(e.response, 'text'):742 error_details += f" - Response: {e.response.text[:200]}"743 print(f"发送消息失败: {error_details}")744 return jsonify({"error": f"Failed to send message: {error_details}"}), 500745 746 747def format_message(messages):748 buffer = io.StringIO()749 role_map, prefix, messages = extract_role(messages)750 for message in messages:751 role = message.get("role")752 role = "\b" + role_map[role] if prefix else role_map[role]753 content = message.get("content").replace("\\n", "\n")754 pattern = re.compile(r"<\|removeRole\|>\n")755 if pattern.match(content):756 content = pattern.sub("", content)757 buffer.write(f"{content}\n")758 else:759 buffer.write(f"{role}: {content}\n\n")760 formatted_message = buffer.getvalue()761 return formatted_message762 763 764def extract_role(messages):765 role_map = {"user": "Human", "assistant": "Assistant", "system": "System"}766 prefix = False767 first_message = messages[0]["content"]768 pattern = re.compile(769 r"""770 <roleInfo>\s*771 user:\s*(?P<user>[^\n]*)\s*772 assistant:\s*(?P<assistant>[^\n]*)\s*773 system:\s*(?P<system>[^\n]*)\s*774 prefix:\s*(?P<prefix>[^\n]*)\s*775 </roleInfo>\n776 """,777 re.VERBOSE,778 )779 match = pattern.search(first_message)780 if match:781 role_map = {782 "user": match.group("user"),783 "assistant": match.group("assistant"),784 "system": match.group("system"),785 }786 prefix = match.group("prefix") == "1"787 messages[0]["content"] = pattern.sub("", first_message)788 print(f"Extracted role map:")789 print(790 f"User: {role_map['user']}, Assistant: {role_map['assistant']}, System: {role_map['system']}"791 )792 print(f"Using prefix: {prefix}")793 return (role_map, prefix, messages)794 795 796@app.route("/health", methods=["GET"])797def health_check():798 global health_check_counter799 health_check_counter += 1800 return jsonify({801 "status": "healthy",802 "timestamp": datetime.now().isoformat(),803 "checks": health_check_counter804 })805 806 807def keep_alive():808 """每20分钟进行一次自我健康检查"""809 while True:810 try:811 requests.get("http://127.0.0.1:7860/health")812 time.sleep(1200) # 20分钟813 except:814 pass # 忽略错误，保持运行815 816 817@app.route("/", methods=["GET"])818def index():819 # 如果需要密码且用户未登录，重定向到登录页面820 if PASSWORD and not flask_session.get('logged_in'):821 return redirect(url_for('login'))822 823 # 否则重定向到仪表盘824 return redirect(url_for('dashboard'))825 826 827# 获取OpenAI的tokenizer来计算token数828def num_tokens_from_string(string, model="gpt-3.5-turbo"):829 """计算文本的token数量"""830 try:831 encoding = tiktoken.encoding_for_model(model)832 num_tokens = len(encoding.encode(string))833 print(f"使用tiktoken计算token数: {num_tokens}")834 return num_tokens835 except Exception as e:836 # 如果tiktoken不支持模型或者出错，使用简单的估算837 estimated_tokens = len(string) // 4 # 粗略估计每个token约4个字符838 print(f"使用估算方法计算token数: {estimated_tokens} (原因: {str(e)})")839 return estimated_tokens840 841# 更新模型使用统计842def update_model_stats(model, prompt_tokens, completion_tokens):843 global model_usage_stats, total_tokens, model_usage_records844 845 # 添加调用记录846 # 获取UTC时间847 utc_now = datetime.utcnow()848 # 转换为北京时间 (UTC+8)849 beijing_time = utc_now + timedelta(hours=8)850 call_time = beijing_time.strftime('%Y-%m-%d %H:%M:%S') # 北京时间851 852 record = {853 "model": model,854 "call_time": call_time,855 "prompt_tokens": prompt_tokens,856 "completion_tokens": completion_tokens,857 "calculation_method": "tiktoken" if any(x in model.lower() for x in ["gpt", "claude"]) or model in ["llama-3", "mistral", "gemma"] else "estimate"858 }859 model_usage_records.append(record)860 861 # 限制记录数量，保留最新的500条862 if len(model_usage_records) > 500:863 model_usage_records.pop(0)864 865 # 保存调用记录到本地文件866 save_model_usage_records()867 868 # 更新聚合统计869 if model not in model_usage_stats:870 model_usage_stats[model] = {871 "count": 0,872 "prompt_tokens": 0,873 "completion_tokens": 0,874 "total_tokens": 0875 }876 877 model_usage_stats[model]["count"] += 1878 model_usage_stats[model]["prompt_tokens"] += prompt_tokens879 model_usage_stats[model]["completion_tokens"] += completion_tokens880 model_usage_stats[model]["total_tokens"] += (prompt_tokens + completion_tokens)881 882 total_tokens["prompt"] += prompt_tokens883 total_tokens["completion"] += completion_tokens884 total_tokens["total"] += (prompt_tokens + completion_tokens)885 886 887# 获取计算点信息888def get_compute_points():889 global compute_points, USER_DATA, users_compute_points890 891 if USER_NUM == 0:892 return893 894 # 清空用户计算点列表895 users_compute_points = []896 897 # 累计总计算点898 total_left = 0899 total_points = 0900 901 # 获取每个用户的计算点信息902 for i, user_data in enumerate(USER_DATA):903 try:904 session, cookies, session_token, _, _ = user_data905 906 # 检查token是否有效907 if is_token_expired(session_token):908 session_token = refresh_token(session, cookies)909 if not session_token:910 print(f"用户{i+1}刷新token失败，无法获取计算点信息")911 continue912 USER_DATA[i] = (session, cookies, session_token, user_data[3], user_data[4])913 914 headers = {915 "accept": "application/json, text/plain, */*",916 "accept-language": "zh-CN,zh;q=0.9",917 "baggage": f"sentry-environment=production,sentry-release=93da8385541a6ce339b1f41b0c94428c70657e22,sentry-public_key=3476ea6df1585dd10e92cdae3a66ff49,sentry-trace_id={TRACE_ID}",918 "reai-ui": "1",919 "sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"",920 "sec-ch-ua-mobile": "?0",921 "sec-ch-ua-platform": "\"Windows\"",922 "sec-fetch-dest": "empty",923 "sec-fetch-mode": "cors",924 "sec-fetch-site": "same-origin",925 "sentry-trace": SENTRY_TRACE,926 "session-token": session_token,927 "x-abacus-org-host": "apps",928 "cookie": cookies929 }930 931 response = session.get(932 COMPUTE_POINTS_URL,933 headers=headers934 )935 936 if response.status_code == 200:937 result = response.json()938 if result.get("success") and "result" in result:939 data = result["result"]940 left = data.get("computePointsLeft", 0)941 total = data.get("totalComputePoints", 0)942 used = total - left943 percentage = round((used / total) * 100, 2) if total > 0 else 0944 945 # 获取北京时间946 beijing_now = datetime.utcnow() + timedelta(hours=8)947 948 # 添加到用户列表949 user_points = {950 "user_id": i + 1, # 用户ID从1开始951 "left": left,952 "total": total,953 "used": used,954 "percentage": percentage,955 "last_update": beijing_now956 }957 users_compute_points.append(user_points)958 959 # 累计总数960 total_left += left961 total_points += total962 963 print(f"用户{i+1}计算点信息更新成功: 剩余 {left}, 总计 {total}")964 965 # 对于第一个用户，获取计算点使用日志966 if i == 0:967 get_compute_points_log(session, cookies, session_token)968 else:969 print(f"获取用户{i+1}计算点信息失败: {result.get('error', '未知错误')}")970 else:971 print(f"获取用户{i+1}计算点信息失败，状态码: {response.status_code}")972 except Exception as e:973 print(f"获取用户{i+1}计算点信息异常: {e}")974 975 # 更新全局计算点信息（所有用户总和）976 if users_compute_points:977 compute_points["left"] = total_left978 compute_points["total"] = total_points979 compute_points["used"] = total_points - total_left980 compute_points["percentage"] = round((compute_points["used"] / compute_points["total"]) * 100, 2) if compute_points["total"] > 0 else 0981 compute_points["last_update"] = datetime.utcnow() + timedelta(hours=8) # 北京时间982 print(f"所有用户计算点总计: 剩余 {total_left}, 总计 {total_points}")983 984# 获取计算点使用日志985def get_compute_points_log(session, cookies, session_token):986 global compute_points_log987 988 try:989 headers = {990 "accept": "application/json, text/plain, */*",991 "accept-language": "zh-CN,zh;q=0.9",992 "content-type": "application/json",993 "reai-ui": "1",994 "sec-ch-ua": "\"Chromium\";v=\"116\", \"Not)A;Brand\";v=\"24\", \"Google Chrome\";v=\"116\"",995 "sec-ch-ua-mobile": "?0",996 "sec-ch-ua-platform": "\"Windows\"",997 "sec-fetch-dest": "empty",998 "sec-fetch-mode": "cors",999 "sec-fetch-site": "same-site",1000 "session-token": session_token,1001 "x-abacus-org-host": "apps",1002 "cookie": cookies1003 }1004 1005 response = session.post(1006 COMPUTE_POINTS_LOG_URL,1007 headers=headers,1008 json={"byLlm": True}1009 )1010 1011 if response.status_code == 200:1012 result = response.json()1013 if result.get("success") and "result" in result:1014 data = result["result"]1015 compute_points_log["columns"] = data.get("columns", {})1016 compute_points_log["log"] = data.get("log", [])1017 print(f"计算点使用日志更新成功，获取到 {len(compute_points_log['log'])} 条记录")1018 else:1019 print(f"获取计算点使用日志失败: {result.get('error', '未知错误')}")1020 else:1021 print(f"获取计算点使用日志失败，状态码: {response.status_code}")1022 except Exception as e:1023 print(f"获取计算点使用日志异常: {e}")1024 1025 1026# 添加登录相关路由1027@app.route("/login", methods=["GET", "POST"])1028def login():1029 error = None1030 if request.method == "POST":1031 password = request.form.get("password")1032 if password and hashlib.sha256(password.encode()).hexdigest() == PASSWORD:1033 flask_session['logged_in'] = True1034 flask_session.permanent = True1035 return redirect(url_for('dashboard'))1036 else:1037 # 密码错误时提示使用环境变量密码1038 error = "密码不正确。请使用设置的环境变量 password 或 password.txt 中的值作为密码和API认证密钥。"1039 1040 # 传递空间URL给模板1041 return render_template('login.html', error=error, space_url=SPACE_URL)1042 1043 1044@app.route("/logout")1045def logout():1046 flask_session.clear()1047 return redirect(url_for('login'))1048 1049 1050@app.route("/dashboard")1051@require_auth1052def dashboard():1053 # 在每次访问仪表盘时更新计算点信息1054 get_compute_points()1055 1056 # 计算运行时间（使用北京时间）1057 beijing_now = datetime.utcnow() + timedelta(hours=8)1058 uptime = beijing_now - START_TIME1059 days = uptime.days1060 hours, remainder = divmod(uptime.seconds, 3600)1061 minutes, seconds = divmod(remainder, 60)1062 1063 if days > 0:1064 uptime_str = f"{days}天 {hours}小时 {minutes}分钟"1065 elif hours > 0:1066 uptime_str = f"{hours}小时 {minutes}分钟"1067 else:1068 uptime_str = f"{minutes}分钟 {seconds}秒"1069 1070 # 当前北京年份1071 beijing_year = beijing_now.year1072 1073 return render_template(1074 'dashboard.html',1075 uptime=uptime_str,1076 health_checks=health_check_counter,1077 user_count=USER_NUM,1078 models=sorted(list(MODELS)),1079 year=beijing_year,1080 model_stats=model_usage_stats,1081 total_tokens=total_tokens,1082 compute_points=compute_points,1083 compute_points_log=compute_points_log,1084 space_url=SPACE_URL, # 传递空间URL1085 users_compute_points=users_compute_points, # 传递用户计算点信息1086 model_usage_records=model_usage_records # 传递模型使用记录1087 )1088 1089 1090# 获取Hugging Face Space URL1091def get_space_url():1092 # 尝试从环境变量获取1093 space_url = os.environ.get("SPACE_URL")1094 if space_url:1095 return space_url1096 1097 # 如果SPACE_URL不存在，尝试从SPACE_ID构建1098 space_id = os.environ.get("SPACE_ID")1099 if space_id:1100 username, space_name = space_id.split("/")1101 # 将空间名称中的下划线替换为连字符1102 # 注意：Hugging Face生成的URL会自动将空间名称中的下划线(_)替换为连字符(-)1103 # 例如："abacus_chat_proxy" 会变成 "abacus-chat-proxy"1104 space_name = space_name.replace("_", "-")1105 return f"https://{username}-{space_name}.hf.space"1106 1107 # 如果以上都不存在，尝试从单独的用户名和空间名构建1108 username = os.environ.get("SPACE_USERNAME")1109 space_name = os.environ.get("SPACE_NAME")1110 if username and space_name:1111 # 将空间名称中的下划线替换为连字符1112 # 同上，Hugging Face会自动进行此转换1113 space_name = space_name.replace("_", "-")1114 return f"https://{username}-{space_name}.hf.space"1115 1116 # 默认返回None1117 return None1118 1119# 获取空间URL1120SPACE_URL = get_space_url()1121if SPACE_URL:1122 print(f"Space URL: {SPACE_URL}")1123 print("注意：Hugging Face生成的URL会自动将空间名称中的下划线(_)替换为连字符(-)")1124 1125 1126if __name__ == "__main__":1127 # 启动保活线程1128 threading.Thread(target=keep_alive, daemon=True).start()1129 1130 # 加载历史模型调用记录1131 load_model_usage_records()1132 1133 # 获取初始计算点信息1134 get_compute_points()1135 1136 port = int(os.environ.get("PORT", 9876))1137 app.run(port=port, host="0.0.0.0")1138