Create app.py

app.py

@@ -0,0 +1,356 @@
+from contextlib import asynccontextmanager
+from datetime import datetime
+from json import dumps, loads
+from logging import 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 playwright.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(INFO)
+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)
+        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 = []
+    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):
+    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('/v1/chat/completions')
+@app.post('/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('/v1/models')
+@app.get('/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('/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('/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('/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: '/v1/chat/completions' },
+                { type: 'locked', path: '/chat/completions' },
+                { type: 'unlocked', path: '/v1/models' },
+                { type: 'unlocked', path: '/models' },
+                { type: 'locked', path: '/update_token' },
+                { type: 'locked', path: '/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)