Improved errors, checking and fixed ratelimit retrying

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())

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',
+        'powered_by': 'nova-api'
-        'by': 'NovaOSS/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'
     })

api/main.py

@@ -35,9 +35,9 @@ limiter = Limiter(
     swallow_errors=True,
     key_func=get_remote_address,
     default_limits=[
-    '1/second',
+    '2/second',
-    '20/minute',
+    '30/minute',
-    '300/hour'
+    '400/hour'
 ])
 
 app.state.limiter = limiter

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:

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...')
+        for func in [test_chat_non_stream_gpt4, test_chat_stream_gpt3]:
-        # print(await test_function_calling())
+            print(f'[*] {func.__name__}')
-
+            result = await func()
-        print('Checking non-streamed chat completions...')
+            print(f'[+] {func.__name__} - {result}')
-        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())
 
     except Exception as exc:
         print('[red]Error: ' + str(exc))