Update app.py

app.py

@@ -2,7 +2,6 @@ from datetime import datetime, timedelta
 from functools import partial
 from os import environ
 from typing import Callable, Coroutine
-from itertools import chain
 
 from anyio import create_task_group
 from fastapi import FastAPI, HTTPException, Request
@@ -11,8 +10,7 @@ from fastapi.responses import StreamingResponse
 from httpx import AsyncClient, RequestError, Timeout
 from starlette.types import Receive, Scope, Send
 
-API_KEYS = ['bad_key'] + [line for line in environ['API_KEYS'].strip().split('\n') if line and line.startswith('sk-')]
-print(f'всего ключей: {len(API_KEYS)}')
+API_KEYS = [line for line in environ['API_KEYS'].strip().split('\n') if line and line.startswith('sk-')]
 COMPLETIONS_URL = 'https://openrouter.ai/api/v1/chat/completions'
 app = FastAPI(title='reverse-proxy')
 
@@ -119,44 +117,32 @@ async def proxy_openai_api(request: Request):
             streaming = client.stream(request.method, COMPLETIONS_URL, headers=headers, params=request.query_params, json=request_body)
             async with streaming as stream_response:
                 if stream_response.status_code in {401, 402, 429}:
-                    print(api_key, 'не валиден')
-                    yield 'auth_error'
-                    return
+                    raise AuthError('ключ API недействителен или превышен лимит отправки запросов')
                 if stream_response.status_code == 403:
-                    raise CensoredError('отклонено по цензуре')
-                
+                    raise CensoredError('отклонено по цензуре')  # это специфичная ошибка именно для опенроутера!
+                response.init_headers({k: v for k, v in stream_response.headers.items() if k not in {'content-length', 'content-encoding', 'alt-svc'}})
+
+                content = bytearray()
                 async for chunk in stream_response.aiter_bytes():
                     if chunk.strip():
                         yield chunk.strip()
+                        content.extend(chunk.strip())
 
         except RequestError as exc:
             raise HTTPException(status_code=500, detail=f'произошла ошибка при запросе: {exc}')
 
-    async def get_response(headers: dict) -> OverrideStreamResponse:
+    async def try_api_keys():
         for api_key in API_KEYS:
-            response_generator = stream_api_response(api_key)
             try:
-                first_chunk = await anext(response_generator)
-                if first_chunk == 'auth_error':
-                    print(f'ключ API {api_key} недействителен или превышен лимит отправки запросов')
-                    continue
-                else:
-                    headers_to_forward = {k: v for k, v in headers.items() if k.lower() not in {'content-length', 'content-encoding', 'alt-svc'}}
-                    
-                    async def combined_generator():
-                        if first_chunk.strip():
-                            yield first_chunk.strip() + b'\n'
-                        async for chunk in response_generator:
-                            if chunk.strip():
-                                yield chunk.strip() + b'\n'
-
-                    return OverrideStreamResponse(combined_generator(), headers=headers_to_forward)
-            except StopAsyncIteration:
+                response_generator = stream_api_response(api_key)
+                response = OverrideStreamResponse(response_generator)
+                return response
+            except AuthError:
+                print(f'ключ API {api_key} недействителен или превышен лимит отправки запросов')
                 continue
         raise HTTPException(status_code=401, detail='все ключи API использованы, доступ запрещен.')
 
-    
-    return await get_response(headers)
+    return await try_api_keys()
 
 
 @cache_results