多token轮询

core/router.py

@@ -23,6 +23,8 @@ logger = setup_logger(__name__)
 router = APIRouter()
 ALLOWED_MODELS = get_settings().ALLOWED_MODELS
 
+current_index = 0
+
 @router.get("/models")
 async def list_models():
     return {"object": "list", "data": ALLOWED_MODELS, "success": True}
@@ -43,56 +45,65 @@ async def chat_completions_options():
 # 文本转语音
 @router.post("/audio/speech")
 async def speech(request: Request):
+    global current_index
     url = 'https://api.thinkbuddy.ai/v1/content/speech/tts'
-    request_headers = {**get_settings().HEADERS,
-                       'authorization': f"Bearer {os.getenv('TOKEN', '')}",
-                       'Accept': 'application/json, text/plain, */*',
-                       }
-    # data = {
-    #     "input": "这是一张插图，显示了一杯饮料，可能是奶昔、冰沙或其他冷饮。杯子上有一个盖子和一根吸管，表明这是一种便于携带和饮用的饮品。这种设计通常用于提供咖啡、冰茶或果汁等饮品。杯子颜色简约，可能用于说明饮品的内容或品牌。",
-    #     "voice": "nova" # alloy echo fable onyx nova shimmer
-    # }
-    body = await request.json()
-    try:
-        async with httpx.AsyncClient(http2=True) as client:
-            response = await client.post(url, headers=request_headers, json=body)
-            response.raise_for_status()
-
-            # 假设响应是音频数据，保存为文件
-            if response.status_code == 200:
-                # 保存音频文件
-                with open('output.mp3', 'wb') as f:
-                    f.write(response.content)
-                print("音频文件已保存为 output.mp3")
-
-                # 异步播放音频
-                # 使用 asyncio.to_thread 来避免阻塞事件循环
-                # await asyncio.to_thread(playsound, 'output.mp3')
-                return True
-            else:
-                print(f"请求失败，状态码: {response.status_code}")
-                print(f"响应内容: {response.text}")
-                return False
-
-    except httpx.RequestError as e:
-        print(f"请求错误: {e}")
-        print("错误堆栈:")
-        traceback.print_exc()
-        return False
-    except httpx.HTTPStatusError as e:
-        print(f"HTTP 错误: {e}")
-        print("错误堆栈:")
-        traceback.print_exc()
-        return False
-    except Exception as e:
-        print(f"发生错误: {e}")
-        print("错误堆栈:")
-        traceback.print_exc()
-        return False
+    token_str = os.getenv('TOKEN', '')
+    token_array = token_str.split(',')
+    if len(token_array) > 0:
+        current_index = current_index % len(token_array)
+        print('speech current index is ', current_index)
+        request_headers = {**get_settings().HEADERS,
+                           'authorization': f"Bearer {token_array[current_index]}",
+                           'Accept': 'application/json, text/plain, */*',
+                           }
+        # data = {
+        #     "input": "这是一张插图，显示了一杯饮料，可能是奶昔、冰沙或其他冷饮。杯子上有一个盖子和一根吸管，表明这是一种便于携带和饮用的饮品。这种设计通常用于提供咖啡、冰茶或果汁等饮品。杯子颜色简约，可能用于说明饮品的内容或品牌。",
+        #     "voice": "nova" # alloy echo fable onyx nova shimmer
+        # }
+        body = await request.json()
+        try:
+            async with httpx.AsyncClient(http2=True) as client:
+                response = await client.post(url, headers=request_headers, json=body)
+                response.raise_for_status()
+
+                # 假设响应是音频数据，保存为文件
+                if response.status_code == 200:
+                    # 保存音频文件
+                    with open('output.mp3', 'wb') as f:
+                        f.write(response.content)
+                    print("音频文件已保存为 output.mp3")
+
+                    # 异步播放音频
+                    # 使用 asyncio.to_thread 来避免阻塞事件循环
+                    # await asyncio.to_thread(playsound, 'output.mp3')
+                    return True
+                else:
+                    print(f"请求失败，状态码: {response.status_code}")
+                    print(f"响应内容: {response.text}")
+                    return False
+
+        except httpx.RequestError as e:
+            print(f"请求错误: {e}")
+            print("错误堆栈:")
+            traceback.print_exc()
+            return False
+        except httpx.HTTPStatusError as e:
+            print(f"HTTP 错误: {e}")
+            print("错误堆栈:")
+            traceback.print_exc()
+            return False
+        except Exception as e:
+            print(f"发生错误: {e}")
+            print("错误堆栈:")
+            traceback.print_exc()
+            return False
+        finally:
+            current_index += 1
 
 # 语音转文本
 @router.post("/audio/transcriptions")
 async def transcriptions(request: Request, file: UploadFile = File(...)):
+    global current_index
     url = 'https://api.thinkbuddy.ai/v1/content/transcribe'
     params = {'enhance': 'true'}
     try:
@@ -112,26 +123,32 @@ async def transcriptions(request: Request, file: UploadFile = File(...)):
         # 记录请求信息
         logger.info(f"Received upload request for file: {file.filename}")
         logger.info(f"Content-Type: {request.headers.get('content-type')}")
-        request_headers = {**get_settings().HEADERS,
-                           'authorization': f"Bearer {os.getenv('TOKEN', '')}",
-                           'Accept': 'application/json, text/plain, */*',
-                           'Content-Type': content_type,
-                           }
-        # 设置较长的超时时间
-        timeout = httpx.Timeout(
-            connect=30.0,  # 连接超时
-            read=300.0,    # 读取超时
-            write=30.0,    # 写入超时
-            pool=30.0      # 连接池超时
-        )
-        # 使用httpx发送异步请求
-        async with httpx.AsyncClient(http2=True, timeout=timeout) as client:
-            response = await client.post(url,
-                                         params=params,
-                                         headers=request_headers,
-                                         files=files)
-            response.raise_for_status()
-            return response.json()
+        token_str = os.getenv('TOKEN', '')
+        token_array = token_str.split(',')
+        if len(token_array) > 0:
+            current_index = current_index % len(token_array)
+            print('transcriptions current index is ', current_index)
+            request_headers = {**get_settings().HEADERS,
+                               'authorization': f"Bearer {token_array[current_index]}",
+                               'Accept': 'application/json, text/plain, */*',
+                               'Content-Type': content_type,
+                               }
+            # 设置较长的超时时间
+            timeout = httpx.Timeout(
+                connect=30.0,  # 连接超时
+                read=300.0,    # 读取超时
+                write=30.0,    # 写入超时
+                pool=30.0      # 连接池超时
+            )
+            # 使用httpx发送异步请求
+            async with httpx.AsyncClient(http2=True, timeout=timeout) as client:
+                response = await client.post(url,
+                                             params=params,
+                                             headers=request_headers,
+                                             files=files)
+                current_index += 1
+                response.raise_for_status()
+                return response.json()
 
     except httpx.TimeoutException:
         raise HTTPException(status_code=504, detail="请求目标服务器超时")
@@ -153,6 +170,7 @@ async def safe_read_file(file: UploadFile) -> Optional[bytes]:
 # 文件上传
 @router.post("/upload")
 async def upload_file(request: Request, file: UploadFile = File(...)):
+    global current_index
     try:
         # 读取文件内容
         content = await safe_read_file(file)
@@ -168,16 +186,22 @@ async def upload_file(request: Request, file: UploadFile = File(...)):
         # 记录请求信息
         logger.info(f"Received upload request for file: {file.filename}")
         logger.info(f"Content-Type: {request.headers.get('content-type')}")
-        request_headers = {**get_settings().HEADERS,
-                           'authorization': f"Bearer {os.getenv('TOKEN', '')}",
-                           'Accept': 'application/json, text/plain, */*',
-                           'Content-Type': content_type,
-                           }
-        # 使用httpx发送异步请求
-        async with httpx.AsyncClient() as client:
-            response = await client.post(f"https://api.thinkbuddy.ai/v1/uploads/images", headers=request_headers,files=files,  timeout=100)
-            response.raise_for_status()
-            return response.json()
+        token_str = os.getenv('TOKEN', '')
+        token_array = token_str.split(',')
+        if len(token_array) > 0:
+            current_index = current_index % len(token_array)
+            print('upload_file current index is ', current_index)
+            request_headers = {**get_settings().HEADERS,
+                               'authorization': f"Bearer {token_array[current_index]}",
+                               'Accept': 'application/json, text/plain, */*',
+                               'Content-Type': content_type,
+                               }
+            # 使用httpx发送异步请求
+            async with httpx.AsyncClient() as client:
+                response = await client.post(f"https://api.thinkbuddy.ai/v1/uploads/images", headers=request_headers,files=files,  timeout=100)
+                current_index += 1
+                response.raise_for_status()
+                return response.json()
 
     except httpx.TimeoutException:
         raise HTTPException(status_code=504, detail="请求目标服务器超时")
@@ -200,6 +224,7 @@ async def upload_file(request: Request, file: UploadFile = File(...)):
 async def chat_completions(
         request: ChatRequest, app_secret: str = Depends(verify_app_secret)
 ):
+    global current_index
     logger.info("Entering chat_completions route")
     # logger.info(f"Received request: {request}")
     # logger.info(f"Received request json format: {json.dumps(request.dict(), indent=4)}")
@@ -216,8 +241,22 @@ async def chat_completions(
 
     if request.stream:
         logger.info("Streaming response")
-        return StreamingResponse(
-            process_streaming_response(request, app_secret),
+
+        # 创建一个标志来追踪是否有响应
+        has_response = False
+
+        async def content_generator():
+            nonlocal has_response
+            try:
+                async for item in process_streaming_response(request, app_secret, current_index):
+                    has_response = True
+                    yield item
+            except Exception as e:
+                logger.error(f"Error in streaming response: {e}")
+                raise
+
+        response = StreamingResponse(
+            content_generator(),
             media_type="text/event-stream",
             headers={
                 "Cache-Control": "no-cache",
@@ -225,6 +264,12 @@ async def chat_completions(
                 "Transfer-Encoding": "chunked"
             }
         )
+
+        # 在返回响应之前增加 current_index
+        # if has_response:
+        #     current_index += 1
+        current_index += 1
+        return response
     else:
         logger.info("Non-streaming response")
         # return await process_non_streaming_response(request)

core/utils.py

@@ -35,32 +35,39 @@ def decode_unicode_escape(s):
 
 FIREBASE_API_KEY = settings.FIREBASE_API_KEY
 async def refresh_token_via_rest(refresh_token):
-    # Firebase Auth REST API endpoint
-    url = f"https://securetoken.googleapis.com/v1/token?key={FIREBASE_API_KEY}"
-
-    payload = {
-        'grant_type': 'refresh_token',
-        'refresh_token': refresh_token
-    }
+    refresh_token_array = [x.strip() for x in refresh_token.split(',')]
+    token_array = []
+    if len(refresh_token_array) > 0:
+        print('refresh token length is ', len(refresh_token_array))
+        for e in refresh_token_array:
+            # Firebase Auth REST API endpoint
+            url = f"https://securetoken.googleapis.com/v1/token?key={FIREBASE_API_KEY}"
+            payload = {
+                'grant_type': 'refresh_token',
+                'refresh_token': e
+            }
 
-    try:
-        async with httpx.AsyncClient() as client:
-            response = await client.post(url, json=payload)
-            if response.status_code == 200:
-                data = response.json()
-                print(json.dumps(data, indent=2))
-                # return {
-                #     'id_token': data['id_token'],
-                #     'refresh_token': data.get('refresh_token'),
-                #     'expires_in': data['expires_in']
-                # }
-                return data['id_token']
-            else:
-                print(f"刷新失败: {response.text}")
+            try:
+                async with httpx.AsyncClient() as client:
+                    response = await client.post(url, json=payload)
+                    if response.status_code == 200:
+                        data = response.json()
+                        print(json.dumps(data, indent=2))
+                        # return {
+                        #     'id_token': data['id_token'],
+                        #     'refresh_token': data.get('refresh_token'),
+                        #     'expires_in': data['expires_in']
+                        # }
+                        # return data['id_token']
+                        token_array.append(data['id_token'])
+                    else:
+                        print(f"刷新失败: {response.text}")
+                        return None
+            except Exception as e:
+                print(f"请求异常: {e}")
                 return None
-    except Exception as e:
-        print(f"请求异常: {e}")
-        return None
+
+        return ','.join(token_array)
 
 
 async def sign_in_with_idp():
@@ -185,7 +192,7 @@ def create_chat_completion_data(
         "usage": None,
     }
 
-async def process_streaming_response(request: ChatRequest, app_secret: str):
+async def process_streaming_response(request: ChatRequest, app_secret: str, current_index: int):
     # 创建自定义 SSL 上下文
     ssl_context = ssl.create_default_context()
     ssl_context.check_hostname = True
@@ -196,57 +203,63 @@ async def process_streaming_response(request: ChatRequest, app_secret: str):
             # http2=True    # 启用 HTTP/2
     ) as client:
         try:
-            request_headers = {**settings.HEADERS, 'authorization': f"Bearer {os.getenv('TOKEN', '')}"} # 从环境变量中获取新的TOKEN
+            token_str = os.getenv('TOKEN', '')
+            token_array = token_str.split(',')
+            if len(token_array) > 0:
+                current_index = current_index % len(token_array)
+                print('completions current index is ', current_index)
+                request_headers = {**settings.HEADERS, 'authorization': f"Bearer {token_array[current_index]}"} # 从环境变量中获取新的TOKEN
+
+                # 直接使用 request.model_dump() 或 request.dict() 获取字典格式的数据
+                request_data = request.model_dump()  # 如果使用较新版本的 Pydantic
+                # # 获取请求数据
+                # request_data = {
+                #     "model": request.model,
+                #     "messages": [msg.dict() for msg in request.messages],
+                #     "temperature": request.temperature,
+                #     "top_p": request.top_p,
+                #     "max_tokens": request.max_tokens,
+                #     "stream": request.stream
+                # }
+                # print("Request Headers:", json.dumps(request_headers, indent=2))  # 格式化打印
+                # print("Request Body:", json.dumps(request.json(), indent=4, ensure_ascii=False))  # 格式化打印
+                print("Request Headers:", json.dumps(request_headers, indent=4, ensure_ascii=False, sort_keys=True, separators=(',', ': ')))  # 格式化打印
+                print("Request Body:", json.dumps(request.model_dump(), indent=4, ensure_ascii=False, sort_keys=True, separators=(',', ': ')))  # 格式化打印
 
-            # 直接使用 request.model_dump() 或 request.dict() 获取字典格式的数据
-            request_data = request.model_dump()  # 如果使用较新版本的 Pydantic
-            # # 获取请求数据
-            # request_data = {
-            #     "model": request.model,
-            #     "messages": [msg.dict() for msg in request.messages],
-            #     "temperature": request.temperature,
-            #     "top_p": request.top_p,
-            #     "max_tokens": request.max_tokens,
-            #     "stream": request.stream
-            # }
-            # print("Request Headers:", json.dumps(request_headers, indent=2))  # 格式化打印
-            # print("Request Body:", json.dumps(request.json(), indent=4, ensure_ascii=False))  # 格式化打印
-            print("Request Headers:", json.dumps(request_headers, indent=4, ensure_ascii=False, sort_keys=True, separators=(',', ': ')))  # 格式化打印
-            print("Request Body:", json.dumps(request.model_dump(), indent=4, ensure_ascii=False, sort_keys=True, separators=(',', ': ')))  # 格式化打印
+                async with client.stream(
+                        "POST",
+                        f"https://api.thinkbuddy.ai/v1/chat/completions",
+                        headers=request_headers,
+                        json=request_data,
+                        timeout=100,
+                ) as response:
+                    response.raise_for_status()
+                    timestamp = int(datetime.now().timestamp())
+                    async for line in response.aiter_lines():
+                        # print(f"{type(line)}: {line}")
+                        if line and line.startswith("data: "):
+                            try:
+                                if line.strip() == 'data: [DONE]':
+                                    await response.aclose()
+                                    break
+                                data_str = line[6:]  # 去掉 'data: ' 前缀
 
-            async with client.stream(
-                    "POST",
-                    f"https://api.thinkbuddy.ai/v1/chat/completions",
-                    headers=request_headers,
-                    json=request_data,
-                    timeout=100,
-            ) as response:
-                response.raise_for_status()
-                timestamp = int(datetime.now().timestamp())
-                async for line in response.aiter_lines():
-                    # print(f"{type(line)}: {line}")
-                    if line and line.startswith("data: "):
-                        try:
-                            if line.strip() == 'data: [DONE]':
-                                await response.aclose()
-                                break
-                            data_str = line[6:]  # 去掉 'data: ' 前缀
+                                # 解析JSON
+                                json_data = json.loads(data_str)
+                                if 'choices' in json_data and len(json_data['choices']) > 0:
+                                    delta = json_data['choices'][0].get('delta', {})
+                                    if 'content' in delta:
+                                        print(delta['content'], end='', flush=True)
+                                        yield f"data: {json.dumps(create_chat_completion_data(delta['content'], request.model, timestamp))}\n\n"
 
-                            # 解析JSON
-                            json_data = json.loads(data_str)
-                            if 'choices' in json_data and len(json_data['choices']) > 0:
-                                delta = json_data['choices'][0].get('delta', {})
-                                if 'content' in delta:
-                                    print(delta['content'], end='', flush=True)
-                                    yield f"data: {json.dumps(create_chat_completion_data(delta['content'], request.model, timestamp))}\n\n"
+                            except json.JSONDecodeError as e:
+                                print(f"JSON解析错误: {e}")
+                                print(f"原始数据: {line}")
+                                continue
 
-                        except json.JSONDecodeError as e:
-                            print(f"JSON解析错误: {e}")
-                            print(f"原始数据: {line}")
-                            continue
+                    yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
+                    yield "data: [DONE]\n\n"
 
-                yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
-                yield "data: [DONE]\n\n"
         except ConnectError as e:
             logger.error(f"Connection error details: {str(e)}")
             raise HTTPException(