app.py

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
from fastapi.responses import HTMLResponse
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)}')
COMPLETIONS_URL = 'https://openrouter.ai/api/v1/chat/completions'
app = FastAPI(title='reverse-proxy')


class Cache:
    def __init__(self, expire: timedelta):
        self.expire = expire
        self.cache = {}
        self.timestamp = datetime.now()

    async def get(self, key):
        if datetime.now() - self.timestamp > self.expire:
            self.cache.clear()
            self.timestamp = datetime.now()
        return self.cache.get(key)

    async def set(self, key, value):
        self.cache[key] = value


cache = Cache(expire=timedelta(hours=1))


def cache_results(func):
    async def wrapper(*args, **kwargs):
        cache_key = f"{func.__name__}:{args}:{kwargs}"
        cached_result = await cache.get(cache_key)
        if cached_result is not None:
            return cached_result
        result = await func(*args, **kwargs)
        await cache.set(cache_key, result)
        return result

    return wrapper


class AuthError(Exception):
    pass


class CensoredError(Exception):
    pass


@app.middleware('http')
async def add_cors_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers['Access-Control-Allow-Origin'] = '*'
    response.headers['Access-Control-Allow-Methods'] = 'GET, POST, PUT, DELETE, PATCH, OPTIONS'
    response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization'
    return response


@app.get('/')
async def root():
    return HTMLResponse('ну пролапс, ну и что')


class OverrideStreamResponse(StreamingResponse):
    async def stream_response(self, send: Send) -> None:
        first_chunk = True
        async for chunk in self.body_iterator:
            if first_chunk:
                await self.send_request_header(send)
                first_chunk = False
            if not isinstance(chunk, bytes):
                chunk = chunk.encode(self.charset)
            await send({'type': 'http.response.body', 'body': chunk, 'more_body': True})

        if first_chunk:
            await self.send_request_header(send)
        await send({'type': 'http.response.body', 'body': b'', 'more_body': False})

    async def send_request_header(self, send: Send) -> None:
        await send({'type': 'http.response.start', 'status': self.status_code, 'headers': self.raw_headers, })

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        async with create_task_group() as task_group:
            async def wrap(func: Callable[[], Coroutine]) -> None:
                await func()
                task_group.cancel_scope.cancel()

            task_group.start_soon(wrap, partial(self.stream_response, send))
            await wrap(partial(self.listen_for_disconnect, receive))

        if self.background is not None:
            await self.background()


async def proxy_openai_api(request: Request):
    headers = {k: v for k, v in request.headers.items() if k not in {'host', 'content-length', 'x-forwarded-for', 'x-real-ip', 'connection'}}

    def update_authorization_header(api_key):
        auth_header_key = next((k for k in headers.keys() if k.lower() == 'authorization'), 'Authorization')
        headers[auth_header_key] = f'Bearer {api_key}'

    client = AsyncClient(verify=False, follow_redirects=True, timeout=Timeout(connect=10, read=90, write=10, pool=10))

    request_body = await request.json() if request.method in {'POST', 'PUT'} else None

    async def stream_api_response(api_key: str):
        update_authorization_header(api_key)
        try:
            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
                if stream_response.status_code == 403:
                    raise CensoredError('отклонено по цензуре')
                
                async for chunk in stream_response.aiter_bytes():
                    if chunk.strip():
                        yield chunk.strip()

        except RequestError as exc:
            raise HTTPException(status_code=500, detail=f'произошла ошибка при запросе: {exc}')

    async def get_response(headers: dict) -> OverrideStreamResponse:
        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:
                continue
        raise HTTPException(status_code=401, detail='все ключи API использованы, доступ запрещен.')

    
    return await get_response(headers)


@cache_results
async def get_free_models():
    async with AsyncClient(follow_redirects=True, timeout=Timeout(10.0, read=30.0, write=10.0, pool=10.0)) as client:
        response = await client.get('https://openrouter.ai/api/v1/models')
        response.raise_for_status()
        data = response.json()
    filtered_models = [model for model in data.get('data', []) if model.get('id', '').endswith(':free')]
    return {'data': filtered_models, 'object': 'list'}


@app.api_route('/v1/models', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
@app.api_route('/models', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
async def get_models():
    return await get_free_models()


@app.api_route('/v1/chat/completions', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
@app.api_route('/chat/completions', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
async def proxy_handler(request: Request):
    return await proxy_openai_api(request)


if __name__ == '__main__':
    from uvicorn import run

    run(app, host='0.0.0.0', port=7860)