app.py

from contextlib import asynccontextmanager
from datetime import datetime
from json import dumps, loads
from logging import DEBUG, Formatter, INFO, StreamHandler, WARNING, getLogger
from os import environ
from pathlib import Path
from random import randint
from typing import AsyncGenerator

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from fastapi import Depends, FastAPI, HTTPException, Header, Request
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, StreamingResponse
from httpx import AsyncClient
from patchright.async_api import Page, async_playwright
from starlette.responses import Response

screenshot_path = Path(__file__).parent / 'screenshot.jpeg'
token_path = Path(__file__).parent / 'token.json'

AUTH_URL = 'https://chat.reka.ai/bff/auth/login'
GET_TOKEN_URL = 'https://chat.reka.ai/bff/auth/access_token'
EMAIL_INPUT = 'input#username'
PASSWRD_INPUT = 'input#password'
SUBMIT_BUTTON = 'button[type="submit"][name="action"][value="default"]:not([data-provider])'
REKA_CHAT_PAGE = 'https://chat.reka.ai/chat'
EMAIL = str(environ.get('EMAIL')).strip()
PASSWRD = str(environ.get('PASSWORD')).strip()
UA = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36 Edg/131.0.0.0'
BROWSER_TIMEOUT_SECONDS = 15

API_TOKEN = str(environ.get('API_TOKEN')).strip()
REKA_API_URL = 'https://chat.reka.ai/api/chat'

logger = getLogger('REKA_API')
logger.setLevel(DEBUG)
handler = StreamHandler()
handler.setLevel(INFO)
formatter = Formatter('%(asctime)s | %(levelname)s : %(message)s', datefmt='%d.%m.%Y %H:%M:%S')
handler.setFormatter(formatter)
logger.addHandler(handler)
getLogger('httpx').setLevel(WARNING)

logger.info('инициализация приложения...')


async def make_screenshot(page: Page):
    await page.screenshot(type='jpeg', path=screenshot_path.resolve().as_posix(), quality=85, full_page=True)
    logger.info('скриншот создан')


async def refresh_token():
    max_timeout = BROWSER_TIMEOUT_SECONDS * 1000
    json_data = {'accessToken': None}
    async with async_playwright() as playwright:
        logger.info('запуск браузера')
        browser = await playwright.chromium.launch(headless=True, args=['--disable-blink-features=AutomationControlled'])
        context = await browser.new_context(
            color_scheme='dark',
            ignore_https_errors=True,
            locale='en-US',
            user_agent=UA,
            no_viewport=True,
        )
        context.set_default_timeout(max_timeout)
        context.set_default_navigation_timeout(max_timeout)
        page = await context.new_page()
        page.set_default_timeout(max_timeout)
        page.set_default_navigation_timeout(max_timeout)
        try:
            logger.info('открытие страницы авторизации')
            await page.goto(AUTH_URL, wait_until='networkidle')
            await make_screenshot(page)
            logger.info('заполнение формы авторизации')
            await page.locator(EMAIL_INPUT).type(EMAIL)
            await page.locator(PASSWRD_INPUT).type(PASSWRD)
            await make_screenshot(page)
            await page.locator(SUBMIT_BUTTON).click()
            logger.info('переход на страницу чата')
            await page.wait_for_url(REKA_CHAT_PAGE)
            await make_screenshot(page)
            logger.info('получение токена')
            await page.goto(GET_TOKEN_URL, wait_until='domcontentloaded')
            await make_screenshot(page)
            json_text = await page.evaluate("document.querySelector('pre').textContent")
            json_data = loads(json_text)
            await make_screenshot(page)
        finally:
            await page.close()
            await context.close()
            await browser.close()
        logger.info('работа браузера завершена')
        return json_data


async def get_token():
    token_data = await refresh_token()
    if token_data.get('accessToken'):
        token_path.write_text(dumps(token_data))
        logger.info('токен reka получен')
    else:
        raise logger.error('токен reka не был получен')


def reka_headers() -> dict[str, str]:
    token = None
    if token_path.exists():
        token = loads(token_path.read_text()).get('accessToken')
    return {
        'authorization': f'Bearer {token}',
        'content-type': 'application/json',
        'user-agent': UA
    }


def convert_openai_to_reka(messages: list[dict]) -> list[dict]:
    reka_messages = [{'type': 'human', 'text': '👋'}]
    skip_next = False
    logger.debug('конвертация сообщений в формат reka')
    for i, message in enumerate(messages):
        if skip_next:
            skip_next = False
            continue
        if message['role'] in ['user', 'assistant']:
            content = message['content']
            if isinstance(content, list):
                text_content = ''
                image_url = None
                for part in content:
                    if part['type'] == 'text':
                        text_content += part['text'] + ' '
                    elif part['type'] == 'image_url':
                        image_url = part['image_url']['url']
                reka_message = {
                    'type': 'human' if message['role'] == 'user' else 'model',
                    'text': text_content.strip()
                }
                if image_url:
                    reka_message['image_url'] = image_url
                    reka_message['media_type'] = 'image'
                reka_messages.append(reka_message)
            else:
                reka_messages.append({
                    'type': 'human' if message['role'] == 'user' else 'model',
                    'text': content
                })
        elif message['role'] == 'system':
            if i + 1 < len(messages) and messages[i + 1]['role'] == 'user':
                combined_text = '[SYSTEM: ' + message['content'] + '] ' + messages[i + 1]['content']
                reka_messages.append({
                    'type': 'human',
                    'text': combined_text
                })
                skip_next = True
            else:
                reka_messages.append({
                    'type': 'human',
                    'text': '[SYSTEM: ' + message['content'] + ']'
                })

    return reka_messages


def format_part(current_text: str, previous_text: str, finish_reason: str):
    logger.debug(f'форматирование сообщения: {current_text}')
    return f"data: {dumps({
        'id': 'chatcmpl-0',
        'object': 'chat.completion.chunk',
        'created': int(datetime.now().timestamp()),
        'model': 'reka-core',
        'system_fingerprint': 'fp_67802d9a6d',
        'choices': [{
            'index': 0,
            'delta': {'content': current_text[len(previous_text):]},
            'finish_reason': finish_reason
        }]}, ensure_ascii=False)}\n\n"


async def fetch_reka_stream(data: dict) -> AsyncGenerator[str, None]:
    logger.info('запрос к reka и стриминг ответа')
    async with AsyncClient() as client:
        response = await client.post(REKA_API_URL, headers=reka_headers(), json=data, timeout=None)
        previous_text = ''
        async for line in response.aiter_lines():
            logger.debug(line)
            if line.startswith('{"detail":'):
                yield format_part('ОШИБКА: ' + loads(line).get('detail'), previous_text, 'error')
                break
            if line.startswith('data:'):
                event_data = loads(line[5:])
                current_text = event_data['text']
                sep_index = current_text.find('<sep')
                finish_reason = None
                if sep_index != -1:
                    finish_reason = 'stop'
                    current_text = current_text[:sep_index].rstrip()
                    if current_text.endswith('\n\n<'):
                        current_text = current_text[:-3]
                if current_text != previous_text:
                    yield format_part(current_text, previous_text, finish_reason)
                    previous_text = current_text
                if finish_reason == 'stop':
                    break


async def periodic_get_token(scheduler: AsyncIOScheduler):
    logger.info('запуск задачи периодического обновления токена')
    scheduler.add_job(
        get_token,
        trigger='interval',
        hours=24,
        next_run_time=datetime.now(),
        misfire_grace_time=3600
    )


@asynccontextmanager
async def app_lifespan(_) -> AsyncGenerator:
    logger.info('запуск приложения')
    scheduler = AsyncIOScheduler()
    await periodic_get_token(scheduler)
    try:
        logger.info('запуск переодических задач')
        scheduler.start()
        logger.info('старт API')
        yield
    finally:
        scheduler.shutdown()
        logger.info('приложение завершено')


app = FastAPI(lifespan=app_lifespan, title='Reka_API')

banned_endpoints = [
    '/openapi.json',
    '/docs',
    '/docs/oauth2-redirect',
    'swagger_ui_redirect',
    '/redoc',
]


@app.middleware('http')
async def block_banned_endpoints(request: Request, call_next):
    logger.debug(f'получен запрос: {request.url.path}')
    if request.url.path in banned_endpoints:
        logger.warning(f'запрещенный endpoint: {request.url.path}')
        return Response(status_code=403)
    response = await call_next(request)
    return response


def verify_token(authorization: str = Header(...)):
    if authorization is None:
        logger.warning('попытка доступа без заголовков авторизации')
        raise HTTPException(status_code=401, detail='эм... нужен пролапс')
    try:
        scheme, token = authorization.split()
        if scheme.lower() != 'bearer':
            logger.warning('попытка доступа с неверным типом авторизации')
            raise HTTPException(status_code=401, detail='пролапс не того вида...')
        if token != API_TOKEN:
            logger.warning('попытка доступа с неверным токеном')
            raise HTTPException(status_code=401, detail='пролапс неверный...')
    except ValueError:
        logger.warning('попытка доступа с неверным типом авторизации')
        raise HTTPException(status_code=401, detail='а где пролапс?')


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


@app.post('/api/chat/completions')
@app.post('/api/v1/chat/completions')
async def chat_completions(request: Request, token: str = Depends(verify_token)):
    logger.debug('запрос `completions`')
    data = await request.json()
    messages = data.get('messages', [])
    reka_messages = convert_openai_to_reka(messages)

    reka_data = {
        'conversation_history': reka_messages,
        'stream': True,
        'use_search_engine': False,
        'use_code_interpreter': False,
        'model_name': 'reka-core',
        'random_seed': randint(0, 2 ** 32 - 1)
    }

    return StreamingResponse(fetch_reka_stream(reka_data), media_type='text/event-stream')


@app.get('/api/models')
@app.get('/api/v1/models')
async def models():
    logger.debug('запрос `models`')
    return JSONResponse({
        'object': 'list',
        'data': [{'id': 'reka-core', 'object': 'model', 'created': int(datetime.now().timestamp()), 'owned_by': 'reka.ai'}]
    }, status_code=200, media_type='application/json')


@app.get('/api/update_token')
async def update_token(token: str = Depends(verify_token)):
    logger.info('запрос `update_token`')
    task = get_token()
    return JSONResponse({'status': 'обновление токена запущено'}, status_code=200, media_type='application/json')


@app.get('/api/show_last_screen')
async def show_last_screen(token: str = Depends(verify_token)):
    logger.info('запрос `show_last_screen`')
    return FileResponse(screenshot_path.resolve().as_posix(), media_type='image/jpeg', status_code=200)

@app.get('/api')
@app.get('/api/info', response_class=HTMLResponse)
async def info():
    logger.debug('запрос `info`')
    return HTMLResponse(content='''<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>Reka Reverse Proxy Endpoints</title>
    <style>
        body {font-family: monospace; background-color: #202020; color:#bfbcb9;}
        .locked::before {content: '🔐'; margin-right: 5px;}
        .unlocked::before {content: '🔓'; margin-right: 5px;}
    </style>
    <script>
        document.addEventListener('DOMContentLoaded', () => {
            const url = `${window.location.protocol}//${window.location.host}`;
            const endpoints = [
                { type: 'locked', path: '/api/v1/chat/completions' },
                { type: 'locked', path: '/api/chat/completions' },
                { type: 'unlocked', path: '/api/v1/models' },
                { type: 'unlocked', path: '/api/models' },
                { type: 'locked', path: '/api/update_token' },
                { type: 'locked', path: '/api/show_last_screen' }
            ];

            const listContainer = document.getElementById('endpoints-list');
            endpoints.forEach(({ type, path }) => {
                const listItem = document.createElement('li');
                listItem.className = type;
                listItem.textContent = `${url}${path}`;
                listContainer.appendChild(listItem);
            });
        });
    </script>
</head>
<body>
    <h2>эндпоинты:</h2>
    <ul id="endpoints-list"></ul>
</body>
</html>''', status_code=200)


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

    logger.info('запуск сервера uvicorn')
    uvicorn_run(app, host='0.0.0.0', port=7860)