gnilets commited on
Commit
899d712
·
verified ·
1 Parent(s): ce0a2f5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -8
app.py CHANGED
@@ -3,7 +3,7 @@ from functools import partial
3
  from os import environ
4
  from typing import Callable, Coroutine
5
  from itertools import chain
6
- from asyncio import anext
7
  from anyio import create_task_group
8
  from fastapi import FastAPI, HTTPException, Request
9
  from fastapi.responses import HTMLResponse
@@ -135,16 +135,17 @@ async def proxy_openai_api(request: Request):
135
  for api_key in API_KEYS:
136
  response_generator = stream_api_response(api_key)
137
  try:
138
- first_chunk = await anext(response_generator)
139
- if first_chunk == 'auth_error':
140
- print(f'ключ API {api_key} недействителен или превышен лимит отправки запросов')
141
- continue
142
- else:
143
- headers_to_forward = {k: v for k, v in headers.items() if k.lower() not in {'content-length', 'content-encoding', 'alt-svc'}}
144
- return OverrideStreamResponse(chain([first_chunk], response_generator), headers=headers_to_forward)
145
  except StopAsyncIteration:
146
  continue
147
  raise HTTPException(status_code=401, detail='все ключи API использованы, доступ запрещен.')
 
148
 
149
  return await get_response()
150
 
 
3
  from os import environ
4
  from typing import Callable, Coroutine
5
  from itertools import chain
6
+
7
  from anyio import create_task_group
8
  from fastapi import FastAPI, HTTPException, Request
9
  from fastapi.responses import HTMLResponse
 
135
  for api_key in API_KEYS:
136
  response_generator = stream_api_response(api_key)
137
  try:
138
+ async for first_chunk in response_generator:
139
+ if first_chunk == 'auth_error':
140
+ print(f'ключ API {api_key} недействителен или превышен лимит отправки запросов')
141
+ break
142
+ else:
143
+ headers_to_forward = {k: v for k, v in headers.items() if k.lower() not in {'content-length', 'content-encoding', 'alt-svc'}}
144
+ return OverrideStreamResponse(chain([first_chunk], response_generator), headers=headers_to_forward)
145
  except StopAsyncIteration:
146
  continue
147
  raise HTTPException(status_code=401, detail='все ключи API использованы, доступ запрещен.')
148
+
149
 
150
  return await get_response()
151