diff --git a/main.py b/main.py index ca8d03f..7763830 100644 --- a/main.py +++ b/main.py @@ -221,7 +221,6 @@ def oaiFreeGetAccessToken(getAccessTokenUrl, refresh_token): 'refresh_token': refresh_token, } response = requests.post(getAccessTokenUrl, data=data) - logging.info(response.text) if not response.ok: logger.error("Request 失败: " + response.text.strip()) return None @@ -336,9 +335,9 @@ scheduler.start() # PANDORA_UPLOAD_URL = 'files.pandoranext.com' -VERSION = '0.7.9.0' +VERSION = '0.7.9.1' # VERSION = 'test' -UPDATE_INFO = '接入oaifree' +UPDATE_INFO = '适配调用team对话,提供查询ChatGPT-Account-ID的/getAccountID接口' # UPDATE_INFO = '【仅供临时测试使用】 ' with app.app_context(): @@ -702,7 +701,7 @@ my_files_types = [ # 定义发送请求的函数 -def send_text_prompt_and_get_response(messages, api_key, stream, model, proxy_api_prefix): +def send_text_prompt_and_get_response(messages, api_key, account_id, stream, model, proxy_api_prefix): url = f"{BASE_URL}{proxy_api_prefix}/backend-api/conversation" headers = { "Authorization": f"Bearer {api_key}" @@ -920,6 +919,11 @@ def send_text_prompt_and_get_response(messages, api_key, stream, model, proxy_ap payload["arkose_token"] = token # 在headers中添加新字段 headers["Openai-Sentinel-Arkose-Token"] = token + + # 用于调用ChatGPT Team次数 + if account_id: + headers["ChatGPT-Account-ID"] = account_id + logger.debug(f"headers: {headers}") logger.debug(f"payload: {payload}") response = requests.post(url, headers=headers, json=payload, stream=True) @@ -1139,7 +1143,8 @@ def replace_sandbox(text, conversation_id, message_id, api_key, proxy_api_prefix return replaced_text -def data_fetcher(upstream_response, data_queue, stop_event, last_data_time, api_key, chat_message_id, model, proxy_api_prefix): +def data_fetcher(upstream_response, data_queue, stop_event, last_data_time, api_key, chat_message_id, model, + proxy_api_prefix): all_new_text = "" first_output = True @@ -1372,7 +1377,7 @@ def data_fetcher(upstream_response, data_queue, stop_event, last_data_time, api_ if is_complete_sandbox_format(file_output_buffer): # 替换完整的引用格式 replaced_text = replace_sandbox(file_output_buffer, conversation_id, - message_id, api_key,proxy_api_prefix) + message_id, api_key, proxy_api_prefix) # print(replaced_text) # 输出替换后的文本 new_text = replaced_text file_output_accumulating = False @@ -1756,13 +1761,18 @@ def chat_completions(): auth_header = request.headers.get('Authorization') if not auth_header or not auth_header.startswith('Bearer '): return jsonify({"error": "Authorization header is missing or invalid"}), 401 - api_key = auth_header.split(' ')[1] + try: + api_key = auth_header.split(' ')[1].split(',')[0].strip() + account_id = auth_header.split(' ')[1].split(',')[1].strip() + logging.info(f"{api_key}:{account_id}") + except IndexError: + account_id = None if not api_key.startswith("eyJhb"): + refresh_token = api_key if api_key in refresh_dict: logger.info(f"从缓存读取到api_key.........。") api_key = refresh_dict.get(api_key) else: - refresh_token = api_key if REFRESH_TOACCESS_ENABLEOAI: api_key = oaiGetAccessToken(api_key) else: @@ -1772,7 +1782,8 @@ def chat_completions(): add_to_dict(refresh_token, api_key) logger.info(f"api_key: {api_key}") - upstream_response = send_text_prompt_and_get_response(messages, api_key, stream, model, proxy_api_prefix) + upstream_response = send_text_prompt_and_get_response(messages, api_key, account_id, stream, model, + proxy_api_prefix) # 在非流式响应的情况下,我们需要一个变量来累积所有的 new_text all_new_text = "" @@ -1791,7 +1802,8 @@ def chat_completions(): # 启动数据处理线程 fetcher_thread = threading.Thread(target=data_fetcher, args=( - upstream_response, data_queue, stop_event, last_data_time, api_key, chat_message_id, model,proxy_api_prefix)) + upstream_response, data_queue, stop_event, last_data_time, api_key, chat_message_id, model, + proxy_api_prefix)) fetcher_thread.start() # 启动保活线程 @@ -1848,9 +1860,9 @@ def chat_completions(): fetcher_thread.join() keep_alive_thread.join() - if conversation_id: - # print(f"准备删除的会话id: {conversation_id}") - delete_conversation(conversation_id, api_key,proxy_api_prefix) + # if conversation_id: + # # print(f"准备删除的会话id: {conversation_id}") + # delete_conversation(conversation_id, api_key,proxy_api_prefix) if not stream: # 执行流式响应的生成函数来累积 all_new_text @@ -1924,14 +1936,19 @@ def images_generations(): auth_header = request.headers.get('Authorization') if not auth_header or not auth_header.startswith('Bearer '): return jsonify({"error": "Authorization header is missing or invalid"}), 401 - api_key = auth_header.split(' ')[1] + try: + api_key = auth_header.split(' ')[1].split(',')[0].strip() + account_id = auth_header.split(' ')[1].split(',')[1].strip() + logging.info(f"{api_key}:{account_id}") + except IndexError: + account_id = None if not api_key.startswith("eyJhb"): + refresh_token = api_key if api_key in refresh_dict: logger.info(f"从缓存读取到api_key.........") api_key = refresh_dict.get(api_key) else: if REFRESH_TOACCESS_ENABLEOAI: - refresh_token = api_key api_key = oaiGetAccessToken(api_key) else: api_key = oaiFreeGetAccessToken(REFRESH_TOACCESS_OAIFREE_REFRESHTOACCESS_URL, api_key) @@ -1951,7 +1968,7 @@ def images_generations(): } ] - upstream_response = send_text_prompt_and_get_response(messages, api_key, False, model,proxy_api_prefix) + upstream_response = send_text_prompt_and_get_response(messages, api_key, account_id, False, model, proxy_api_prefix) # 在非流式响应的情况下,我们需要一个变量来累积所有的 new_text all_new_text = "" @@ -2367,6 +2384,56 @@ def get_file(filename): return send_from_directory('files', filename) +@app.route(f'/{API_PREFIX}/getAccountID' if API_PREFIX else '/getAccountID', methods=['POST']) +@cross_origin() # 使用装饰器来允许跨域请求 +def getAccountID(): + logger.info(f"New Img Request") + proxy_api_prefix = getPROXY_API_PREFIX(lock) + if proxy_api_prefix == None: + return jsonify({"error": "PROXY_API_PREFIX is not accessible"}), 401 + auth_header = request.headers.get('Authorization') + if not auth_header or not auth_header.startswith('Bearer '): + return jsonify({"error": "Authorization header is missing or invalid"}), 401 + try: + api_key = auth_header.split(' ')[1].split(',')[0].strip() + account_id = auth_header.split(' ')[1].split(',')[1].strip() + except IndexError: + account_id = None + if not api_key.startswith("eyJhb"): + refresh_token = api_key + if api_key in refresh_dict: + logger.info(f"从缓存读取到api_key.........") + api_key = refresh_dict.get(api_key) + else: + if REFRESH_TOACCESS_ENABLEOAI: + api_key = oaiGetAccessToken(api_key) + else: + api_key = oaiFreeGetAccessToken(REFRESH_TOACCESS_OAIFREE_REFRESHTOACCESS_URL, api_key) + if not api_key.startswith("eyJhb"): + return jsonify({"error": "refresh_token is wrong or refresh_token url is wrong!"}), 401 + add_to_dict(refresh_token, api_key) + logger.info(f"api_key: {api_key}") + + url = f"{BASE_URL}{proxy_api_prefix}/backend-api/accounts/check/v4-2023-04-27" + headers = { + "Authorization": "Bearer " + api_key + } + res = requests.get(url, headers=headers) + if res.status_code == 200: + data = res.json() + result = {"plus": set(), "team": set()} + for account_id, account_data in data["accounts"].items(): + plan_type = account_data["account"]["plan_type"] + if plan_type == "team": + result[plan_type].add(account_id) + elif plan_type == "plus": + result[plan_type].add(account_id) + result = {plan_type: list(ids) for plan_type, ids in result.items()} + return jsonify(result) + else: + return jsonify({"error": "Request failed."}), 400 + + # 内置自动刷新access_token def updateRefresh_dict(): success_num = 0