diff --git a/api/db/providerkeys.py b/api/db/providerkeys.py index 84930de..4e11253 100644 --- a/api/db/providerkeys.py +++ b/api/db/providerkeys.py @@ -12,7 +12,7 @@ class KeyManager: self.conn = AsyncIOMotorClient(os.environ['MONGO_URI']) async def _get_collection(self, collection_name: str): - return self.conn[os.getenv('MONGO_NAME', 'nova-test')][collection_name] + return self.conn['nova-core'][collection_name] async def add_key(self, provider: str, key: str, source: str='?'): db = await self._get_collection('providerkeys') @@ -36,7 +36,7 @@ class KeyManager: }) if key is None: - return ValueError('No keys available for this provider!') + return '--NO_KEY--' return key['key'] @@ -87,4 +87,4 @@ class KeyManager: manager = KeyManager() if __name__ == '__main__': - asyncio.run(manager.delete_empty_keys()) + asyncio.run(manager.import_all()) diff --git a/api/helpers/errors.py b/api/helpers/errors.py index a4bee00..04fa82d 100644 --- a/api/helpers/errors.py +++ b/api/helpers/errors.py @@ -8,8 +8,7 @@ async def error(code: int, message: str, tip: str) -> starlette.responses.Respon 'code': code, 'message': message, 'tip': tip, - 'website': 'https://nova-oss.com', - 'by': 'NovaOSS/Nova-API' + 'powered_by': 'nova-api' }} return starlette.responses.Response(status_code=code, content=json.dumps(info)) @@ -20,5 +19,6 @@ async def yield_error(code: int, message: str, tip: str) -> str: return json.dumps({ 'code': code, 'message': message, - 'tip': tip + 'tip': tip, + 'powered_by': 'nova-api' }) diff --git a/api/main.py b/api/main.py index f224550..a97d2d4 100644 --- a/api/main.py +++ b/api/main.py @@ -35,9 +35,9 @@ limiter = Limiter( swallow_errors=True, key_func=get_remote_address, default_limits=[ - '1/second', - '20/minute', - '300/hour' + '2/second', + '30/minute', + '400/hour' ]) app.state.limiter = limiter diff --git a/api/responder.py b/api/responder.py index 657a223..8907104 100644 --- a/api/responder.py +++ b/api/responder.py @@ -38,7 +38,6 @@ async def respond( is_chat = False model = None - is_stream = False if 'chat/completions' in path: is_chat = True @@ -73,6 +72,13 @@ async def respond( provider_name = provider_auth.split('>')[0] provider_key = provider_auth.split('>')[1] + if provider_key == '--NO_KEY--': + yield await errors.yield_error(500, + 'Sorry, our API seems to have issues connecting to our provider(s).', + 'This most likely isn\'t your fault. Please try again later.' + ) + return + target_request['headers'].update(target_request.get('headers', {})) if target_request['method'] == 'GET' and not payload: @@ -91,12 +97,13 @@ async def respond( timeout=aiohttp.ClientTimeout( connect=1.0, total=float(os.getenv('TRANSFER_TIMEOUT', '500')) - ), + ) ) as response: is_stream = response.content_type == 'text/event-stream' if response.status == 429: - await keymanager.rate_limit_key(provider_name, provider_key) + print('[!] rate limit') + # await keymanager.rate_limit_key(provider_name, provider_key) continue if response.content_type == 'application/json': @@ -112,12 +119,14 @@ async def respond( critical_error = True if critical_error: + print('[!] critical error') continue if response.ok: server_json_response = client_json_response else: + print('[!] non-ok response', client_json_response) continue if is_stream: @@ -136,10 +145,6 @@ async def respond( except Exception as exc: print('[!] exception', exc) - if 'too many requests' in str(exc): - #!TODO - pass - continue else: diff --git a/checks/client.py b/checks/client.py index f194d23..23252b2 100644 --- a/checks/client.py +++ b/checks/client.py @@ -26,6 +26,12 @@ MESSAGES = [ api_endpoint = os.getenv('CHECKS_ENDPOINT', 'http://localhost:2332/v1') +async def _response_base_check(response: httpx.Response) -> None: + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise ConnectionError(f'API returned an error: {response.json()}') from exc + async def test_server(): """Tests if the API server is running.""" @@ -36,7 +42,7 @@ async def test_server(): url=f'{api_endpoint.replace("/v1", "")}', timeout=3 ) - response.raise_for_status() + await _response_base_check(response) assert response.json()['ping'] == 'pong', 'The API did not return a correct response.' except httpx.ConnectError as exc: @@ -63,7 +69,7 @@ async def test_chat_non_stream_gpt4() -> float: json=json_data, timeout=10, ) - response.raise_for_status() + await _response_base_check(response) assert '1337' in response.json()['choices'][0]['message']['content'], 'The API did not return a correct response.' return time.perf_counter() - request_start @@ -86,7 +92,7 @@ async def test_chat_stream_gpt3() -> float: json=json_data, timeout=10, ) - response.raise_for_status() + await _response_base_check(response) chunks = [] resulting_text = '' @@ -128,7 +134,7 @@ async def test_image_generation() -> float: json=json_data, timeout=10, ) - response.raise_for_status() + await _response_base_check(response) assert '://' in response.json()['data'][0]['url'] return time.perf_counter() - request_start @@ -166,7 +172,7 @@ async def test_function_calling(): json=json_data, timeout=15, ) - response.raise_for_status() + await _response_base_check(response) res = response.json() output = json.loads(res['choices'][0]['message']['function_call']['arguments']) @@ -185,7 +191,7 @@ async def test_models(): headers=HEADERS, timeout=3 ) - response.raise_for_status() + await _response_base_check(response) res = response.json() all_models = [model['id'] for model in res['data']] @@ -208,20 +214,10 @@ async def demo(): else: raise ConnectionError('API Server is not running.') - # print('[lightblue]Checking if function calling works...') - # print(await test_function_calling()) - - print('Checking non-streamed chat completions...') - print(await test_chat_non_stream_gpt4()) - - print('Checking streamed chat completions...') - print(await test_chat_stream_gpt3()) - - # print('[lightblue]Checking if image generation works...') - # print(await test_image_generation()) - - # print('Checking the models endpoint...') - # print(await test_models()) + for func in [test_chat_non_stream_gpt4, test_chat_stream_gpt3]: + print(f'[*] {func.__name__}') + result = await func() + print(f'[+] {func.__name__} - {result}') except Exception as exc: print('[red]Error: ' + str(exc))