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 | |
| import asyncio | |
| import httpx | |
| from concurrent.futures import ThreadPoolExecutor | |
| import atexit | |
| 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) | |
| # 创建一个线程池用于执行 token 计算 | |
| token_executor = ThreadPoolExecutor(max_workers=5) | |
| # 确保程序退出时关闭线程池 | |
| atexit.register(token_executor.shutdown, wait=True) | |
| # 添加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() | |
| 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 | |
| async def refresh_token(client: httpx.AsyncClient, 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: | |
| # 使用 httpx.AsyncClient 发送异步请求 | |
| response = await client.post( | |
| USER_INFO_URL, | |
| headers=headers, | |
| json={}, | |
| # cookies=None # httpx doesn't use cookies param like this with client | |
| ) | |
| 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 | |
| async def get_model_map(client: httpx.AsyncClient, 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: | |
| # 使用 httpx.AsyncClient 发送异步请求 | |
| response = await client.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 | |
| async def init_session(): | |
| get_password() | |
| global USER_NUM, MODELS, USER_DATA | |
| config_list = resolve_config() | |
| user_num = len(config_list) | |
| all_models = set() | |
| # 创建一个列表来收集异步任务 | |
| init_tasks = [] | |
| for i in range(user_num): | |
| user = config_list[i] | |
| cookies = user.get("cookies") | |
| conversation_id = user.get("conversation_id") | |
| # 创建 httpx.AsyncClient 实例 | |
| client = httpx.AsyncClient(timeout=30.0) # 设置超时 | |
| # 将每个用户的初始化变成一个异步任务 | |
| async def initialize_user(i, client, cookies, conversation_id): | |
| session_token = await refresh_token(client, cookies) # await 调用异步函数 | |
| if not session_token: | |
| print(f"无法获取cookie {i+1}的token") | |
| await client.aclose() # 关闭 client | |
| return None | |
| try: | |
| model_map, models_set = await get_model_map(client, cookies, session_token) # await 调用异步函数 | |
| # 用户数据现在包含 httpx.AsyncClient | |
| user_entry = (client, 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 = await client.post( # await 调用异步函数 | |
| 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}") | |
| return user_entry, models_set # 返回用户数据和模型集合 | |
| except Exception as e: | |
| print(f"配置用户 {i+1} 失败: {e}") | |
| await client.aclose() # 关闭 client | |
| return None | |
| # 添加任务到列表 | |
| init_tasks.append(initialize_user(i, client, cookies, conversation_id)) | |
| # 并发执行所有用户初始化任务 | |
| results = await asyncio.gather(*init_tasks) | |
| # 处理结果 | |
| temp_user_data = [] | |
| for result in results: | |
| if result: | |
| user_entry, models_set = result | |
| temp_user_data.append(user_entry) | |
| all_models.update(models_set) | |
| # 按原始索引排序 (虽然 gather 可能保持顺序，但显式排序更安全) | |
| USER_DATA = sorted(temp_user_data, key=lambda x: x[5]) | |
| USER_NUM = len(USER_DATA) | |
| if USER_NUM == 0: | |
| print("No user available, exiting...") | |
| # 在异步函数中不能直接 exit(1)，可以引发异常或返回特殊值 | |
| raise RuntimeError("No user available") | |
| 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 | |
| # 修改 get_user_data 以处理异步 refresh_token | |
| async def get_user_data(): | |
| global CURRENT_USER | |
| CURRENT_USER = (CURRENT_USER + 1) % USER_NUM | |
| print(f"使用配置 {CURRENT_USER+1}") | |
| # Akiru uzantajn datumojn | |
| client, cookies, session_token, conversation_id, model_map, user_index = USER_DATA[CURRENT_USER] | |
| # Kontrolu ĉu la tokeno eksvalidiĝis, se jes, refreŝigu ĝin | |
| if is_token_expired(session_token): | |
| print(f"Cookie {CURRENT_USER+1}的token已过期或即将过期，正在刷新...") | |
| new_token = await refresh_token(client, cookies) | |
| if new_token: | |
| # Ĝisdatigu la globale konservitan tokenon | |
| USER_DATA[CURRENT_USER] = (client, cookies, new_token, conversation_id, model_map, user_index) | |
| session_token = new_token | |
| print(f"成功更新token: {session_token[:15]}...{session_token[-15:]}") | |
| else: | |
| print(f"警告：无法刷新Cookie {CURRENT_USER+1}的token，继续使用当前token") | |
| return (client, 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 - 改为异步提交 | |
| print("\n----- 提交输入token计算任务 -----") | |
| prompt_token_future = token_executor.submit(num_tokens_from_string, message, model) | |
| print("输入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" | |
| # 处理API返回的流式内容 | |
| 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}") | |
| # 结束标记 | |
| yield "data: " + json.dumps({"object": "chat.completion.chunk", "choices": [{"delta": {}, "finish_reason": "stop"}]}) + "\n\n" | |
| yield "data: [DONE]\n\n" | |
| # 内容处理完毕后进行操作（其实不会执行到这里，因为响应已经流式传输给客户端） | |
| try: | |
| # 获取完整响应内容并计算输出Token | |
| completion_text = completion_buffer.getvalue() | |
| # 异步提交输出Token计算 | |
| completion_token_future = token_executor.submit(num_tokens_from_string, completion_text, model) | |
| print("输出token计算已提交到后台") | |
| # 保存对话历史并获取计算点数 | |
| _, compute_points_used = save_conversation_history(session, cookies, session_token, conversation_id) | |
| # 更新统计信息 | |
| update_model_stats(model, prompt_token_future, completion_token_future, compute_points_used) | |
| except Exception as post_e: | |
| print(f"流结束后处理出错: {post_e}") | |
| # 返回流式响应 | |
| 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) | |
| # 保存到配置文件 | |
| try: | |
| update_conversation_id(user_index, new_conversation_id) | |
| except Exception as config_e: | |
| print(f"更新配置文件失败但继续重试: {config_e}") | |
| # 修改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() | |
| # 重新定义extract_segment函数（解决作用域问题） | |
| def retry_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 "" | |
| buffer = io.StringIO() | |
| for line in response.iter_lines(): | |
| if line: | |
| decoded_line = line.decode("utf-8") | |
| segment = retry_extract_segment(decoded_line) | |
| if segment: | |
| buffer.write(segment) | |
| response_content = buffer.getvalue() | |
| # 异步计算输出token | |
| completion_token_future = token_executor.submit(num_tokens_from_string, response_content, model) | |
| print("重试后的输出token计算已提交到后台") | |
| # 保存对话历史并获取计算点数 | |
| _, compute_points_used = save_conversation_history(session, cookies, session_token, new_conversation_id) | |
| # 更新统计信息 | |
| update_model_stats(model, prompt_token_future, completion_token_future, compute_points_used) | |
| # 获取最终结果用于返回 | |
| try: | |
| prompt_tokens = prompt_token_future.result(timeout=10) | |
| if isinstance(prompt_tokens, tuple): | |
| prompt_tokens = prompt_tokens[0] # 提取token数值 | |
| completion_tokens = completion_token_future.result(timeout=10) | |
| if isinstance(completion_tokens, tuple): | |
| completion_tokens = completion_tokens[0] # 提取token数值 | |
| total_tokens = prompt_tokens + completion_tokens | |
| except Exception as token_e: | |
| print(f"获取token结果失败但继续返回: {token_e}") | |
| prompt_tokens = 0 | |
| completion_tokens = 0 | |
| total_tokens = 0 | |
| 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": total_tokens | |
| } | |
| }) | |
| except Exception as retry_e: | |
| print(f"重试失败: {retry_e}") | |
| # 取消token计算任务，避免资源浪费 | |
| try: | |
| if 'prompt_token_future' in locals() and not prompt_token_future.done(): | |
| prompt_token_future.cancel() | |
| except Exception as cancel_e: | |
| print(f"取消token计算任务失败: {cancel_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() | |
| # 计算输入token - 改为异步提交 | |
| print("\n----- 提交输入token计算任务 -----") | |
| prompt_token_future = token_executor.submit(num_tokens_from_string, message, model) | |
| print("输入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 "" | |
| 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 | |
| full_content = f"<think>\n{think_content}\n</think>\n{response_content}" | |
| completion_token_future = token_executor.submit(num_tokens_from_string, full_content, model) | |
| print("输出token计算已提交到后台") | |
| # 保存对话历史并获取计算点数 | |
| _, compute_points_used = save_conversation_history(session, cookies, session_token, conversation_id) | |
| # 更新统计信息 | |
| update_model_stats(model, prompt_token_future, completion_token_future, compute_points_used) | |
| # 获取最终结果用于返回 | |
| try: | |
| prompt_tokens = prompt_token_future.result(timeout=10) | |
| if isinstance(prompt_tokens, tuple): | |
| prompt_tokens = prompt_tokens[0] # 提取token数值 | |
| completion_tokens = completion_token_future.result(timeout=10) | |
| if isinstance(completion_tokens, tuple): | |
| completion_tokens = completion_tokens[0] # 提取token数值 | |
| total_tokens = prompt_tokens + completion_tokens | |
| except Exception as token_e: | |
| print(f"获取token结果失败但继续返回: {token_e}") | |
| prompt_tokens = 0 | |
| completion_tokens = 0 | |
| total_tokens = 0 | |
| 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": total_tokens | |
| } | |
| }) | |
| else: | |
| 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 | |
| completion_token_future = token_executor.submit(num_tokens_from_string, response_content, model) | |
| print("输出token计算已提交到后台") | |
| # 保存对话历史并获取计算点数 | |
| _, compute_points_used = save_conversation_history(session, cookies, session_token, conversation_id) | |
| # 更新统计信息 | |
| update_model_stats(model, prompt_token_future, completion_token_future, compute_points_used) | |
| # 获取最终结果用于返回 | |
| try: | |
| prompt_tokens = prompt_token_future.result(timeout=10) | |
| if isinstance(prompt_tokens, tuple): | |
| prompt_tokens = prompt_tokens[0] # 提取token数值 | |
| completion_tokens = completion_token_future.result(timeout=10) | |
| if isinstance(completion_tokens, tuple): | |
| completion_tokens = completion_tokens[0] # 提取token数值 | |
| total_tokens = prompt_tokens + completion_tokens | |
| except Exception as token_e: | |
| print(f"获取token结果失败但继续返回: {token_e}") | |
| prompt_tokens = 0 | |
| completion_tokens = 0 | |
| total_tokens = 0 | |
| 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": total_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) | |
| # 保存到配置文件 | |
| try: | |
| update_conversation_id(user_index, new_conversation_id) | |
| except Exception as config_e: | |
| print(f"更新配置文件失败但继续重试: {config_e}") | |
| # 修改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() | |
| # 重新定义extract_segment函数（解决作用域问题） | |
| def retry_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 "" | |
| buffer = io.StringIO() | |
| for line in response.iter_lines(): | |
| if line: | |
| decoded_line = line.decode("utf-8") | |
| segment = retry_extract_segment(decoded_line) | |
| if segment: | |
| buffer.write(segment) | |
| response_content = buffer.getvalue() | |
| # 异步计算输出token | |
| completion_token_future = token_executor.submit(num_tokens_from_string, response_content, model) | |
| print("重试后的输出token计算已提交到后台") | |
| # 保存对话历史并获取计算点数 | |
| _, compute_points_used = save_conversation_history(session, cookies, session_token, new_conversation_id) | |
| # 更新统计信息 | |
| update_model_stats(model, prompt_token_future, completion_token_future, compute_points_used) | |
| # 获取最终结果用于返回 | |
| try: | |
| prompt_tokens = prompt_token_future.result(timeout=10) | |
| if isinstance(prompt_tokens, tuple): | |
| prompt_tokens = prompt_tokens[0] # 提取token数值 | |
| completion_tokens = completion_token_future.result(timeout=10) | |
| if isinstance(completion_tokens, tuple): | |
| completion_tokens = completion_tokens[0] # 提取token数值 | |
| total_tokens = prompt_tokens + completion_tokens | |
| except Exception as token_e: | |
| print(f"获取token结果失败但继续返回: {token_e}") | |
| prompt_tokens = 0 | |
| completion_tokens = 0 | |
| total_tokens = 0 | |
| 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": total_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_token_future, completion_token_future, compute_points=None): | |
| global model_usage_stats, total_tokens, model_usage_records | |
| prompt_tokens = 0 | |
| prompt_calc_method = "估算(获取失败)" | |
| try: | |
| # 在这里等待并获取输入 token 计算结果，设置超时 | |
| print("等待输入 token 计算结果...") | |
| prompt_tokens, prompt_calc_method = prompt_token_future.result(timeout=20) # 超时时间设长一点 | |
| print(f"输入 token 结果: {prompt_tokens} ({prompt_calc_method})") | |
| except TimeoutError: | |
| print("获取输入 token 计算结果超时") | |
| prompt_calc_method = "估算(超时)" | |
| # 估算 - 如果收到了原始输入，可以在这里估算 | |
| except Exception as e: | |
| print(f"获取输入 token 计算结果失败: {e}") | |
| # 保持 0 和失败标记 | |
| completion_tokens = 0 | |
| completion_calc_method = "估算(获取失败)" | |
| try: | |
| # 在这里等待并获取输出 token 计算结果，设置超时 | |
| print("等待输出 token 计算结果...") | |
| completion_tokens, completion_calc_method = completion_token_future.result(timeout=20) # 超时时间设长一点 | |
| print(f"输出 token 结果: {completion_tokens} ({completion_calc_method})") | |
| except TimeoutError: | |
| print("获取输出 token 计算结果超时") | |
| completion_calc_method = "估算(超时)" | |
| # 可以选择估算或保持 0 | |
| except Exception as e: | |
| print(f"获取输出 token 计算结果失败: {e}") | |
| # 保持 0 和失败标记 | |
| # 添加调用记录 | |
| # 获取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": f"输入: {prompt_calc_method}, 输出: {completion_calc_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) | |
| print("模型统计信息已更新") | |
| # 获取计算点信息 | |
| 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") | |