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/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/logs.py

@@ -1,21 +1,17 @@
 import os
 import time
+import asyncio
 
 from dotenv import load_dotenv
 from motor.motor_asyncio import AsyncIOMotorClient
 
-from helpers import network
+try:
+    from helpers import network
+except ImportError:
+    pass
 
 load_dotenv()
 
-UA_SIMPLIFY = {
-    'Windows NT': 'W',
-    'Mozilla/5.0': 'M',
-    'Win64; x64': '64',
-    'Safari/537.36': 'S',
-    'AppleWebKit/537.36 (KHTML, like Gecko)': 'K',
-}
-
 ## MONGODB Setup
 
 conn = AsyncIOMotorClient(os.environ['MONGO_URI'])
@@ -30,18 +26,7 @@ async def replacer(text: str, dict_: dict) -> str:
     return text
 
 async def log_api_request(user: dict, incoming_request, target_url: str):
-    """Logs the API Request into the database.
-    No input prompt is logged, however data such as IP & useragent is noted.
-    This would be useful for security reasons. Other minor data is also collected.
-
-    Args:
-        user (dict): User dict object
-        incoming_request (_type_): Request
-        target_url (str): The URL the api request was targetted to.
-
-    Returns:
-        _type_: _description_
-    """
+    """Logs the API Request into the database."""
     db = await _get_collection('logs')
     payload = {}
 
@@ -53,7 +38,6 @@ async def log_api_request(user: dict, incoming_request, target_url: str):
 
     model = payload.get('model')
     ip_address = await network.get_ip(incoming_request)
-    useragent = await replacer(incoming_request.headers.get('User-Agent', ''), UA_SIMPLIFY)
 
     new_log_item = {
         'timestamp': time.time(),
@@ -62,7 +46,6 @@ async def log_api_request(user: dict, incoming_request, target_url: str):
         'user_id': str(user['_id']),
         'security': {
             'ip': ip_address,
-            'useragent': useragent,
         },
         'details': {
             'model': model,
@@ -90,5 +73,21 @@ async def delete_by_user_id(user_id: str):
     db = await _get_collection('logs')
     return await db.delete_many({'user_id': user_id})
 
+async def get_logs_time_range(start: int, end: int):
+    db = await _get_collection('logs')
+
+    entries = []
+    async for entry in db.find({'timestamp': {'$gte': start, '$lte': end}}):
+        entries.append(entry)
+
+    return entries
+
+async def main():
+    # how many requests in last 24 hours?
+    last_24_hours = time.time() - 86400
+    logs = await get_logs_time_range(last_24_hours, time.time())
+
+    print(f'Number of logs in last 24 hours: {len(logs)}')
+
 if __name__ == '__main__':
-    pass
+    asyncio.run(main())

api/db/providerkeys.py

@@ -3,6 +3,8 @@ import time
 import random
 import asyncio
 
+import aiofiles
+import aiofiles.os
 from aiocache import cached
 from dotenv import load_dotenv
 from cachetools import TTLCache
@@ -72,10 +74,10 @@ 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
 
@@ -98,7 +100,7 @@ manager = KeyManager()
 
 async def main():
     keys = await manager.get_possible_keys('closed')
-    print(len(keys))
+    print(keys)
 
 if __name__ == '__main__':
     asyncio.run(main())

api/db/stats.py

@@ -2,8 +2,6 @@ import os
 import pytz
 import asyncio
 import datetime
-import json
-import time
 
 from dotenv import load_dotenv
 from motor.motor_asyncio import AsyncIOMotorClient
@@ -15,13 +13,6 @@ load_dotenv()
 class StatsManager:
     """
     ### The manager for all statistics tracking
-    Stats tracked:
-    - Dates
-    - IPs
-    - Target URLs
-    - Tokens
-    - Models
-    - URL Paths
     """
 
     def __init__(self):

api/db/tester.py

@@ -1,6 +0,0 @@
-from stats import * 
-import asyncio
-
-manager = StatsManager()
-
-asyncio.run(manager.get_model_usage())

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

api/providers/__init__.py

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

api/providers/__main__.py

@@ -3,14 +3,13 @@ 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)
@@ -22,7 +21,7 @@ async def main():
     except IndexError:
         print('List of available providers:')
 
-        for file_name in os.listdir(os.path.dirname(__file__)):
+        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])
 

api/providers/azure.py

@@ -12,7 +12,7 @@ MODELS = [
     'gpt-4',
     'gpt-4-32k'
 ]
-MODELS = [f'{model}-azure' for model in MODELS]
+# MODELS = [f'{model}-azure' for model in MODELS]
 
 AZURE_API = '2023-07-01-preview'
 

api/providers/helpers/utils.py

@@ -34,7 +34,4 @@ async def conversation_to_prompt(conversation: list) -> str:
     return text
 
 async def random_secret_for(name: str) -> str:
-    try:
-        return await providerkeys.manager.get_key(name)
-    except ValueError:
-        raise ValueError(f'Keys missing for "{name}" <no_keys>')
+    return await providerkeys.manager.get_key(name)

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

@@ -2,7 +2,7 @@
 
 import os
 import json
-import logging
+import ujson
 import aiohttp
 import asyncio
 import starlette
@@ -49,9 +49,7 @@ async def respond(
         'Content-Type': 'application/json'
     }
 
-    for i in range(20):
-        print(i)
-        # Load balancing: randomly selecting a suitable provider
+    for i in range(1):
         try:
             if is_chat:
                 target_request = await load_balancing.balance_chat_request(payload)
@@ -151,13 +149,21 @@ async def respond(
 
                         async for chunk in response.content.iter_any():
                             chunk = chunk.decode('utf8').strip()
+
+                            if 'azure' in provider_name:
+                                chunk = chunk.strip().replace('data: ', '')
+
+                                if not chunk or 'prompt_filter_results' in chunk:
+                                    continue
+
                             yield chunk + '\n\n'
 
                     break
 
             except Exception as exc:
                 print('[!] exception', exc)
-                continue
+                # continue
+                raise exc
 
     else:
         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.')

checks/client.py

@@ -215,7 +215,10 @@ async def demo():
         else:
             raise ConnectionError('API Server is not running.')
 
-        for func in [test_chat_non_stream_gpt4, test_chat_stream_gpt3]:
+        for func in [
+            # test_chat_non_stream_gpt4,
+            test_chat_stream_gpt3
+        ]:
             print(f'[*] {func.__name__}')
             result = await func()
             print(f'[+] {func.__name__} - {result}')

playground/notes.txt

@@ -0,0 +1,10 @@
+--- EXPECTED ---
+
+data: {"id":"custom-chatcmpl-nUSiapqELukaPT7vEnGcXkbvrS1fR","object":"chat.completion.chunk","created":1696716717,"model":"gpt-3.5-turbo-0613","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}
+
+data: {"id":"custom-chatcmpl-nUSiapqELukaPT7vEnGcXkbvrS1fR","object":"chat.completion.chunk","created":1696716717,"model":"gpt-3.5-turbo-0613","choices":[{"index":0,"delta":{"content":"123"},"finish_reason":null}]}
+
+data: {"id":"custom-chatcmpl-nUSiapqELukaPT7vEnGcXkbvrS1fR","object":"chat.completion.chunk","created":1696716717,"model":"gpt-3.5-turbo-0613","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
+
+data: [DONE]
+

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__':