Refactor file operations

api/backup_manager/main.py

@@ -1,6 +1,8 @@
 import os
 import json
 import asyncio
+import aiofiles
+import aiofiles.os
 
 from sys import argv
 from bson import json_util
@@ -18,8 +20,7 @@ async def main(output_dir: str):
 async def make_backup(output_dir: str):
     output_dir = os.path.join(FILE_DIR, '..', 'backups', output_dir)
 
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
+    await aiofiles.os.makedirs(output_dir, exist_ok=True)
 
     client = AsyncIOMotorClient(MONGO_URI)
     databases = await client.list_database_names()
@@ -29,22 +30,22 @@ async def make_backup(output_dir: str):
         if database == 'local':
             continue
 
-        if not os.path.exists(f'{output_dir}/{database}'):
-            os.mkdir(f'{output_dir}/{database}')
+        await aiofiles.os.makedirs(os.path.join(output_dir, database), exist_ok=True)
 
         for collection in databases[database]:
             print(f'Initiated database backup for {database}/{collection}')
             await make_backup_for_collection(database, collection, output_dir)
 
 async def make_backup_for_collection(database, collection, output_dir):
-    path = f'{output_dir}/{database}/{collection}.json'
+    path = os.path.join(output_dir, database, f'{collection}.json')
 
     client = AsyncIOMotorClient(MONGO_URI)
     collection = client[database][collection]
     documents = await collection.find({}).to_list(length=None)
 
-    with open(path, 'w') as f:
-        json.dump(documents, f, default=json_util.default)
+    async with aiofiles.open(path, 'w') as f:
+        for chunk in json.JSONEncoder(default=json_util.default).iterencode(documents):
+            await f.write(chunk)
 
 if __name__ == '__main__':
     if len(argv) < 2 or len(argv) > 2:

api/config/config.yml

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

api/core.py

@@ -13,6 +13,7 @@ import json
 import hmac
 import httpx
 import fastapi
+import aiofiles
 import functools
 
 from dhooks import Webhook, Embed
@@ -148,11 +149,14 @@ async def run_checks(incoming_request: fastapi.Request):
 async def get_crypto_price(cryptocurrency: str) -> float:
     """Gets the price of a cryptocurrency using coinbase's API."""
 
-    if os.path.exists('cache/crypto_prices.json'):
-        with open('cache/crypto_prices.json', 'r') as f:
-            cache = json.load(f)
-    else:
+    cache_path = os.path.join('cache', 'crypto_prices.json')
+    try:
+        async with aiofiles.open(cache_path) as f:
+            content = await f.read()
+    except FileNotFoundError:
         cache = {}
+    else:
+        cache = json.loads(content)
 
     is_old = time.time() - cache.get('_last_updated', 0) > 60 * 60
 
@@ -164,8 +168,9 @@ async def get_crypto_price(cryptocurrency: str) -> float:
             cache[cryptocurrency] = usd_price
             cache['_last_updated'] = time.time()
 
-            with open('cache/crypto_prices.json', 'w') as f:
-                json.dump(cache, f)
+            async with aiofiles.open(cache_path, 'w') as f:
+                for chunk in json.JSONEncoder().iterencode(cache):
+                    await f.write(chunk)
 
     return cache[cryptocurrency]
 

api/db/providerkeys.py

@@ -1,12 +1,19 @@
 import os
 import time
+import random
 import asyncio
 
+import aiofiles
+import aiofiles.os
+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'])
@@ -24,27 +31,34 @@ class KeyManager:
             'source': source,
         })
 
-    async def get_key(self, provider: str):
+    async def get_possible_keys(self, provider: str):
         db = await self._get_collection('providerkeys')
-        key = await db.find_one({
+        keys = await db.find({
             'provider': provider,
             'inactive_reason': None,
             '$or': [
-                {'rate_limited_since': None},
-                {'rate_limited_since': {'$lte': time.time() - 86400}}
+                {'rate_limited_until': None},
+                {'rate_limited_until': {'$lte': time.time()}}
             ]
-        })
+        }).to_list(length=None)
 
-        if key is None:
+        return keys
+
+    async def get_key(self, provider: str):
+        keys = await self.get_possible_keys(provider)
+
+        if not keys:
             return '--NO_KEY--'
 
-        return key['key']
+        key = random.choice(keys)
+        api_key = key['key']
+        return api_key
 
-    async def rate_limit_key(self, provider: str, key: str):
+    async def rate_limit_key(self, provider: str, key: str, duration: int):
         db = await self._get_collection('providerkeys')
         await db.update_one({'provider': provider, 'key': key}, {
             '$set': {
-                'rate_limited_since': time.time()
+                'rate_limited_until': time.time() + duration
             }
         })
 
@@ -60,18 +74,16 @@ class KeyManager:
         db = await self._get_collection('providerkeys')
         num = 0
 
-        for filename in os.listdir('api/secret'):
+        for filename in await aiofiles.os.listdir(os.path.join('api', 'secret')):
             if filename.endswith('.txt'):
-                with open(f'api/secret/{filename}') as f:
-                    for line in f.readlines():
+                async with aiofiles.open(os.path.join('api', 'secret', 'filename')) as f:
+                    async for line in f:
                         if not line.strip():
                             continue
 
                         await db.insert_one({
                             'provider': filename.split('.')[0],
                             'key': line.strip(),
-                            'rate_limited_since': None,
-                            'inactive_reason': None,
                             'source': 'import'
                         })
                         num += 1
@@ -86,5 +98,9 @@ class KeyManager:
 
 manager = KeyManager()
 
+async def main():
+    keys = await manager.get_possible_keys('closed')
+    print(len(keys))
+
 if __name__ == '__main__':
-    asyncio.run(manager.import_all())
+    asyncio.run(main())

api/db/users.py

@@ -14,7 +14,7 @@ except ImportError:
 
 load_dotenv()
 
-with open(helpers.root + '/api/config/config.yml', encoding='utf8') as f:
+with open(os.path.join(helpers.root, 'api', 'config', 'config.yml'), encoding='utf8') as f:
     credits_config = yaml.safe_load(f)
 
 ## MONGODB Setup

api/handler.py

@@ -19,10 +19,11 @@ from helpers import tokens, errors, network
 load_dotenv()
 
 users = UserManager()
-models_list = json.load(open('cache/models.json', encoding='utf8'))
+with open(os.path.join('cache', 'models.json'), encoding='utf8') as f:
+    models_list = json.load(f)
 models = [model['id'] for model in models_list['data']]
 
-with open('config/config.yml', encoding='utf8') as f:
+with open(os.path.join('config', 'config.yml'), encoding='utf8') as f:
     config = yaml.safe_load(f)
 
 moderation_debug_key_key = os.getenv('MODERATION_DEBUG_KEY')
@@ -69,8 +70,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('NO_RATELIMIT_KEYS')
-    if '/enterprise' in path and user.get('api_key') not in enterprise_keys:
+    enterprise_keys = os.environ.get('ENTERPRISE_KEYS')
+    if path.startswith('/enterprise/v1') and user.get('api_key') not in enterprise_keys.split():
         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,9 +1,11 @@
 from . import \
+    azure, \
     closed, \
     closed4
     # closed432
 
 MODULES = [
+    azure,
     closed,
     closed4,
     # closed432,

api/providers/__main__.py

@@ -1,52 +1,56 @@
 import os
 import sys
+import aiohttp
+import asyncio
+import importlib
+import aiofiles.os
 
 from rich import print
 
 def remove_duplicate_keys(file):
     with open(file, 'r', encoding='utf8') as f:
-        lines = f.readlines()
-
-    unique_lines = set(lines)
+        unique_lines = set(f)
 
     with open(file, 'w', encoding='utf8') as f:
         f.writelines(unique_lines)
 
-try:
-    provider_name = sys.argv[1]
+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}')
+    except IndexError:
+        print('List of available providers:')
 
-        exit()
+        for file_name in await aiofiles.os.listdir(os.path.dirname(__file__)):
+            if file_name.endswith('.py') and not file_name.startswith('_'):
+                print(file_name.split('.')[0])
 
-except IndexError:
-    print('List of available providers:')
+        sys.exit(0)
 
-    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])
+    try:
+        provider = importlib.import_module(f'.{provider_name}', 'providers')
+    except ModuleNotFoundError as exc:
+        print(exc)
+        sys.exit(1)
 
-    sys.exit(0)
+    if len(sys.argv) > 2:
+        model = sys.argv[2] # choose a specific model
+    else:
+        model = provider.MODELS[-1] # choose best model
 
-try:
-    provider = __import__(provider_name)
-except ModuleNotFoundError as 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)
+    print(f'{provider_name} @ {model}')
+    req = await provider.chat_completion(model=model, messages=[{'role': 'user', 'content': '1+1='}])
+    print(req)
 
-if len(sys.argv) > 2:
-    model = sys.argv[2]
-else:
-    model = provider.MODELS[-1]
+    # 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)
 
-
-print(f'{provider_name} @ {model}')
-comp = provider.chat_completion(model=model)
-print(comp)
+asyncio.run(main())

api/providers/azure.py

@@ -0,0 +1,32 @@
+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,4 +1,7 @@
-from db import providerkeys
+try:
+    from db import providerkeys
+except ModuleNotFoundError:
+    from ...db import providerkeys
 
 GPT_3 = [
     'gpt-3.5-turbo',

api/providers/mandrill.py

@@ -0,0 +1,21 @@
+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/proxies.py

@@ -96,7 +96,7 @@ proxies_in_files = []
 
 for proxy_type in ['http', 'socks4', 'socks5']:
     try:
-        with open(f'secret/proxies/{proxy_type}.txt') as f:
+        with open(os.path.join('secret', 'proxies', f'{proxy_type}.txt')) as f:
             for line in f:
                 clean_line = line.split('#', 1)[0].strip()
                 if clean_line:

api/responder.py

@@ -49,7 +49,8 @@ async def respond(
         'Content-Type': 'application/json'
     }
 
-    for _ in range(20):
+    for i in range(20):
+        print(i)
         # Load balancing: randomly selecting a suitable provider
         try:
             if is_chat:
@@ -62,6 +63,7 @@ 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
@@ -73,6 +75,7 @@ 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.'
@@ -101,16 +104,26 @@ 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()
 
-                        if 'method_not_supported' in str(client_json_response):
-                            await errors.error(500, 'Sorry, this endpoint does not support this method.', data['error']['message'])
+                        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
 
                         critical_error = False
                         for error in CRITICAL_API_ERRORS:
@@ -126,7 +139,6 @@ 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,6 +25,7 @@ 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:

requirements.txt

@@ -1,3 +1,4 @@
+aiofiles==23.2.1
 aiohttp==3.8.5
 aiohttp_socks==0.8.0
 dhooks==1.1.4

rewards/main.py

@@ -51,7 +51,7 @@ async def update_roles():
 def launch():
     asyncio.run(main())
 
-    with open('rewards/last_update.txt', 'w', encoding='utf8') as f:
+    with open(os.path.join('rewards', 'last_update.txt'), 'w', encoding='utf8') as f:
         f.write(str(time.time()))
 
 if __name__ == '__main__':