Fix dangling asyncio tasks

api/config/config.yml

@@ -6,10 +6,9 @@ costs:
   other: 5
 
   chat-models:
-    gpt-4-32k-azure: 100
-    gpt-4: 50
-    gpt-4-azure: 10
-    gpt-3: 5
+    gpt-4-32k: 100
+    gpt-4: 30
+    gpt-3: 3
 
 ## Roles Explanation
 

api/db/providerkeys.py

@@ -1,17 +1,12 @@
 import os
 import time
-import random
 import asyncio
 
-from aiocache import cached
 from dotenv import load_dotenv
-from cachetools import TTLCache
 from motor.motor_asyncio import AsyncIOMotorClient
 
 load_dotenv()
 
-cache = TTLCache(maxsize=100, ttl=10)
-
 class KeyManager:
     def __init__(self):
         self.conn = AsyncIOMotorClient(os.environ['MONGO_URI'])
@@ -29,34 +24,27 @@ class KeyManager:
             'source': source,
         })
 
-    async def get_possible_keys(self, provider: str):
+    async def get_key(self, provider: str):
         db = await self._get_collection('providerkeys')
-        keys = await db.find({
+        key = await db.find_one({
             'provider': provider,
             'inactive_reason': None,
             '$or': [
-                {'rate_limited_until': None},
-                {'rate_limited_until': {'$lte': time.time()}}
+                {'rate_limited_since': None},
+                {'rate_limited_since': {'$lte': time.time() - 86400}}
             ]
-        }).to_list(length=None)
+        })
 
-        return keys
-
-    async def get_key(self, provider: str):
-        keys = await self.get_possible_keys(provider)
-
-        if not keys:
+        if key is None:
             return '--NO_KEY--'
 
-        key = random.choice(keys)
-        api_key = key['key']
-        return api_key
+        return key['key']
 
-    async def rate_limit_key(self, provider: str, key: str, duration: int):
+    async def rate_limit_key(self, provider: str, key: str):
         db = await self._get_collection('providerkeys')
         await db.update_one({'provider': provider, 'key': key}, {
             '$set': {
-                'rate_limited_until': time.time() + duration
+                'rate_limited_since': time.time()
             }
         })
 
@@ -82,6 +70,8 @@ class KeyManager:
                         await db.insert_one({
                             'provider': filename.split('.')[0],
                             'key': line.strip(),
+                            'rate_limited_since': None,
+                            'inactive_reason': None,
                             'source': 'import'
                         })
                         num += 1
@@ -96,9 +86,5 @@ class KeyManager:
 
 manager = KeyManager()
 
-async def main():
-    keys = await manager.get_possible_keys('closed')
-    print(len(keys))
-
 if __name__ == '__main__':
-    asyncio.run(main())
+    asyncio.run(manager.import_all())

api/handler.py

@@ -69,8 +69,8 @@ async def handle(incoming_request: fastapi.Request):
         return await errors.error(403, f'Your NovaAI account has been banned. Reason: \'{ban_reason}\'.', 'Contact the staff for an appeal.')
 
     # Checking for enterprise status
-    enterprise_keys = os.environ.get('ENTERPRISE_KEYS')
-    if path.startswith('/enterprise/v1') and user.get('api_key') not in enterprise_keys.split():
+    enterprise_keys = os.environ.get('NO_RATELIMIT_KEYS')
+    if '/enterprise' in path and user.get('api_key') not in enterprise_keys:
         return await errors.error(403, 'Enterprise API is not available.', 'Contact the staff for an upgrade.')
 
     if 'account/credits' in path:

api/providers/__init__.py

@@ -1,11 +1,9 @@
 from . import \
-    azure, \
     closed, \
     closed4
     # closed432
 
 MODULES = [
-    azure,
     closed,
     closed4,
     # closed432,

api/providers/__main__.py

@@ -1,8 +1,5 @@
 import os
 import sys
-import aiohttp
-import asyncio
-import importlib
 
 from rich import print
 
@@ -15,10 +12,16 @@ def remove_duplicate_keys(file):
     with open(file, 'w', encoding='utf8') as f:
         f.writelines(unique_lines)
 
-async def main():
 try:
     provider_name = sys.argv[1]
 
+    if provider_name == '--clear':
+        for file in os.listdir('secret/'):
+            if file.endswith('.txt'):
+                remove_duplicate_keys(f'secret/{file}')
+
+        exit()
+
 except IndexError:
     print('List of available providers:')
 
@@ -29,29 +32,21 @@ async def main():
     sys.exit(0)
 
 try:
-        provider = importlib.import_module(f'.{provider_name}', 'providers')
+    provider = __import__(provider_name)
 except ModuleNotFoundError as exc:
-        print(exc)
+    print(f'Provider "{provider_name}" not found.')
+    print('Available providers:')
+    for file_name in os.listdir(os.path.dirname(__file__)):
+        if file_name.endswith('.py') and not file_name.startswith('_'):
+            print(file_name.split('.')[0])
     sys.exit(1)
 
 if len(sys.argv) > 2:
-        model = sys.argv[2] # choose a specific model
+    model = sys.argv[2]
 else:
-        model = provider.MODELS[-1] # choose best model
+    model = provider.MODELS[-1]
+
 
 print(f'{provider_name} @ {model}')
-    req = await provider.chat_completion(model=model, messages=[{'role': 'user', 'content': '1+1='}])
-    print(req)
-
-    # launch aiohttp
-    async with aiohttp.ClientSession() as session:
-        async with session.request(
-            method=req['method'],
-            url=req['url'],
-            headers=req['headers'],
-            json=req['payload'],
-        ) as response:
-            res_json = await response.json()
-            print(response.status, res_json)
-
-asyncio.run(main())
+comp = provider.chat_completion(model=model)
+print(comp)

api/providers/azure.py

@@ -1,32 +0,0 @@
-from .helpers import utils
-
-AUTH = True
-ORGANIC = False
-CONTEXT = True
-STREAMING = True
-MODERATIONS = False
-ENDPOINT = 'https://nova-00001.openai.azure.com'
-MODELS = [
-    'gpt-3.5-turbo',
-    'gpt-3.5-turbo-16k',
-    'gpt-4',
-    'gpt-4-32k'
-]
-MODELS = [f'{model}-azure' for model in MODELS]
-
-AZURE_API = '2023-07-01-preview'
-
-async def chat_completion(**payload):
-    key = await utils.random_secret_for('azure-nva1')
-
-    deployment = payload['model'].replace('.', '').replace('-azure', '')
-
-    return {
-        'method': 'POST',
-        'url': f'{ENDPOINT}/openai/deployments/{deployment}/chat/completions?api-version={AZURE_API}',
-        'payload': payload,
-        'headers': {
-            'api-key': key
-        },
-        'provider_auth': f'azure-nva1>{key}'
-    }

api/providers/helpers/utils.py

@@ -1,7 +1,4 @@
-try:
 from db import providerkeys
-except ModuleNotFoundError:
-    from ...db import providerkeys
 
 GPT_3 = [
     'gpt-3.5-turbo',

api/providers/mandrill.py

@@ -1,21 +0,0 @@
-from .helpers import utils
-
-AUTH = True
-ORGANIC = False
-CONTEXT = True
-STREAMING = True
-MODELS = ['llama-2-7b-chat']
-
-async def chat_completion(**kwargs):
-    payload = kwargs
-    key = await utils.random_secret_for('mandrill')
-
-    return {
-        'method': 'POST',
-        'url': f'https://api.mandrillai.tech/v1/chat/completions',
-        'payload': payload,
-        'headers': {
-            'Authorization': f'Bearer {key}'
-        },
-        'provider_auth': f'mandrill>{key}'
-    }

api/responder.py

@@ -52,8 +52,7 @@ async def respond(
         'Content-Type': 'application/json'
     }
 
-    for i in range(20):
-        print(i)
+    for _ in range(20):
         # Load balancing: randomly selecting a suitable provider
         try:
             if is_chat:
@@ -66,7 +65,6 @@ async def respond(
                     'headers': headers,
                     'cookies': incoming_request.cookies
                 })
-
         except ValueError:
             yield await errors.yield_error(500, f'Sorry, the API has no active API keys for {model}.', 'Please use a different model.')
             return
@@ -78,7 +76,6 @@ async def respond(
             provider_key = provider_auth.split('>')[1]
 
         if provider_key == '--NO_KEY--':
-            print(f'No key for {provider_name}')
             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.'
@@ -107,26 +104,16 @@ async def respond(
                 ) as response:
                     is_stream = response.content_type == 'text/event-stream'
 
+                    if response.status == 429:
+                        print('[!] rate limit')
+                        # await keymanager.rate_limit_key(provider_name, provider_key)
+                        continue
+
                     if response.content_type == 'application/json':
                         client_json_response = await response.json()
 
-                        try:
-                            error_code = client_json_response['error']['code']
-                        except KeyError:
-                            error_code = ''
-
-                        if error_code == 'method_not_supported':
-                            yield await errors.yield_error(400, 'Sorry, this endpoint does not support this method.', 'Please use a different method.')
-
-                        if error_code == 'insufficient_quota':
-                            print('[!] insufficient quota')
-                            await keymanager.rate_limit_key(provider_name, provider_key, 86400)
-                            continue
-
-                        if error_code == 'billing_not_active':
-                            print('[!] billing not active')
-                            await keymanager.deactivate_key(provider_name, provider_key, 'billing_not_active')
-                            continue
+                        if 'method_not_supported' in str(client_json_response):
+                            await errors.error(500, 'Sorry, this endpoint does not support this method.', data['error']['message'])
 
                         critical_error = False
                         for error in CRITICAL_API_ERRORS:
@@ -142,6 +129,7 @@ async def respond(
                             server_json_response = client_json_response
 
                         else:
+                            print('[!] non-ok response', client_json_response)
                             continue
 
                     if is_stream:

checks/client.py

@@ -25,7 +25,6 @@ MESSAGES = [
 ]
 
 api_endpoint = os.getenv('CHECKS_ENDPOINT', 'http://localhost:2332/v1')
-# api_endpoint = 'http://localhost:2333/v1'
 
 async def _response_base_check(response: httpx.Response) -> None:
     try: