Various improvements

README.md

@@ -1,6 +1,12 @@
 # ☄️ NovaOSS API Server
 Reverse proxy server for "Closed"AI's API.
 
+> "*OpenAI is very closed*"
+> 
+> — [ArsTechnica (July 2023)](https://arstechnica.com/information-technology/2023/07/is-chatgpt-getting-worse-over-time-study-claims-yes-but-others-arent-sure/)
+
+We aim to fix that!
+
 ## NovaOSS APIs
 Our infrastructure might seem a bit confusing, but it's actually quite simple. Just the first one really matters for you, if you want to access our AI API. The other ones are just for the team.
 

api/closed_keys.py

@@ -1,114 +0,0 @@
-"""ClosedAI key manager."""
-
-import os
-import uuid
-import time
-import asyncio
-
-from rich import print
-from helpers import databases
-
-async def prepare() -> None:
-    """Creates the database tables"""
-
-    keys_db = await databases.connect('closed_keys')
-    await keys_db.execute(
-        """
-        CREATE TABLE IF NOT EXISTS closed_keys (
-            id TEXT PRIMARY KEY,
-            key TEXT,
-            source TEXT DEFAULT 'unknown',
-            created_at INTEGER,
-            last_used INTEGER,
-            uses_count INTEGER DEFAULT 0,
-            tokens_generated INTEGER DEFAULT 0,
-            active BOOLEAN,
-            working BOOLEAN,
-            tags TEXT DEFAULT ''
-        )
-        """
-    )
-    await keys_db.commit()
-
-async def add_key(
-    key: str,
-    source: str='unknown',
-    tags: list=None
-) -> None:
-    """Adds a new key to the database"""
-
-    tags = tags or []
-
-    new_key = {
-        'id': str(uuid.uuid4()),
-        'key': key,
-        'source': source,
-        'created_at': int(time.time()),
-        'last_used': -1,
-        'uses_count': 0,
-        'tokens_generated': 0,
-        'active': True,
-        'working': True,
-        'tags': '/'.join(tags),
-    }
-
-    await databases.insert_dict(new_key, 'closed_keys', 'closed_keys')
-
-async def get_working_key() -> dict:
-    """Returns a working key"""
-
-    keys_db = await databases.connect('closed_keys')
-
-    async with keys_db.execute('SELECT * FROM closed_keys WHERE working = 1') as cursor:
-        async for row in cursor:
-            return await databases.row_to_dict(row, cursor)
-
-    return None
-
-asyncio.run(prepare())
-
-async def key_stopped_working(key: str) -> None:
-    """Marks a key as stopped working"""
-
-    keys_db = await databases.connect('closed_keys')
-
-    await keys_db.execute(
-        """
-        UPDATE closed_keys
-        SET working = 0
-        WHERE key = :key
-        """,
-        {'key': key}
-    )
-    await keys_db.commit()
-
-async def key_was_used(key: str, num_tokens: int) -> None:
-    """Updates the stats of a key"""
-
-    keys_db = await databases.connect('closed_keys')
-
-    # set last_used to int of time.time(), adds one to uses_count and adds num_tokens to tokens_generated
-    await keys_db.execute(
-        """
-        UPDATE closed_keys
-        SET last_used = :last_used, uses_count = uses_count + 1, tokens_generated = tokens_generated + :tokens_generated
-        WHERE key = :key
-        """,
-        {
-            'key': key,
-            'last_used': int(time.time()),
-            'tokens_generated': num_tokens
-        }
-    )
-    await keys_db.commit()
-
-async def demo():
-    await add_key('sk-non-working-key')
-    await key_stopped_working('sk-non-working-key')
-    await add_key('sk-working-key')
-    key = await get_working_key()
-    print(key)
-
-if __name__ == '__main__':
-    asyncio.run(demo())
-    os.system(f'pkill -f {os.path.basename(__file__)}')

api/helpers/databases.py

@@ -1,36 +0,0 @@
-"""Database helper."""
-
-import os
-import aiosqlite
-
-from dotenv import load_dotenv
-
-load_dotenv()
-
-async def row_to_dict(row, cursor):
-    """Converts a database row to into a <dict>."""    
-
-    return {col[0]: row[idx] for idx, col in enumerate(cursor.description)}
-
-async def connect(name: str):
-    """Creates a connection to the database"""
-
-    return await aiosqlite.connect(f'{os.getenv("API_DB_PATH")}{name}.db')
-
-async def insert_dict(dict_: dict, table: str, db) -> dict:
-    """Adds a dictionary to a table, safely."""
-
-    if isinstance(db, str):
-        db = await connect(db)
-
-    sep = ', '
-    query = f"""
-        INSERT INTO {table} ({sep.join(dict_.keys())})
-        VALUES ({sep.join([f':{key}' for key in dict_.keys()])})
-    """
-
-    await db.execute(query, dict_)
-    await db.commit()
-    
-    return dict_
-

api/helpers/errors.py

@@ -0,0 +1,13 @@
+import json
+import starlette
+
+def error(code: int, message: str, tip: str) -> starlette.responses.Response: 
+    info = {'error': {
+        'code': code,
+        'message': message,
+        'tip': tip,
+        'website': 'https://nova-oss.com',
+        'by': 'NovaOSS/Nova-API'
+    }}
+
+    return starlette.responses.Response(status_code=code, content=json.dumps(info))

api/helpers/tokens.py

@@ -0,0 +1,48 @@
+import tiktoken
+
+def count_for_messages(messages: list, model: str='gpt-3.5-turbo-0613') -> int:
+    """Return the number of tokens used by a list of messages."""
+
+    try:
+        encoding = tiktoken.encoding_for_model(model)
+
+    except KeyError:
+        encoding = tiktoken.get_encoding('cl100k_base')
+
+    if model in {
+        'gpt-3.5-turbo-0613',
+        'gpt-3.5-turbo-16k-0613',
+        'gpt-4-0314',
+        'gpt-4-32k-0314',
+        'gpt-4-0613',
+        'gpt-4-32k-0613',
+    }:
+        tokens_per_message = 3
+        tokens_per_name = 1
+
+    elif model == 'gpt-3.5-turbo-0301':
+        tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
+        tokens_per_name = -1  # if there's a name, the role is omitted
+
+    elif 'gpt-3.5-turbo' in model:
+        return num_tokens_from_messages(messages, model='gpt-3.5-turbo-0613')
+
+    elif 'gpt-4' in model:
+        return num_tokens_from_messages(messages, model='gpt-4-0613')
+
+    else:
+        raise NotImplementedError(f"""num_tokens_from_messages() is not implemented for model {model}.
+See https://github.com/openai/openai-python/blob/main/chatml.md
+for information on how messages are converted to tokens.""")
+    
+    num_tokens = 0
+    for message in messages:
+        num_tokens += tokens_per_message
+        for key, value in message.items():
+            num_tokens += len(encoding.encode(value))
+            if key == 'name':
+                num_tokens += tokens_per_name
+
+    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
+
+    return num_tokens

api/main.py

@@ -39,4 +39,4 @@ async def root():
         'github': 'https://github.com/novaoss/nova-api'
     }
 
-app.add_route('/v1/{path:path}', transfer.handle_api_request, ['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])
+app.add_route('/{path:path}', transfer.handle_api_request, ['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])

api/users.py

@@ -1,5 +1,3 @@
-"""User system."""
-
 import os
 import uuid
 import time
@@ -7,38 +5,32 @@ import string
 import random
 import asyncio
 
-from helpers import databases
 from dotenv import load_dotenv
+from motor.motor_asyncio import AsyncIOMotorClient
+
+from rich import print
 
 load_dotenv()
 
+MONGO_URI = os.getenv('MONGO_URI')
+MONGO_DB_NAME = 'users'
+
+def get_mongo(collection_name):
+    client = AsyncIOMotorClient(MONGO_URI)
+    db = client[MONGO_DB_NAME]
+    return db[collection_name]
+
 async def prepare() -> None:
-    """Creates the database tables"""
+    """Create the MongoDB collection."""
 
-    users_db = await databases.connect('users')
+    collection = get_mongo('users')
-    await users_db.execute(
-        """
-        CREATE TABLE IF NOT EXISTS users (
-            id TEXT PRIMARY KEY,
-            api_key TEXT,
-            active BOOLEAN,
-            created_at INTEGER,
-            last_used INTEGER DEFAULT 0,
-            uses_count INTEGER DEFAULT 0,
-            tokens_generated INTEGER DEFAULT 0,
-            discord_id INTEGER DEFAULT 0,
-            credit INTEGER DEFAULT 0,
-            tags TEXT DEFAULT ''
-        )
-        """
-    )
-    await users_db.commit()
 
-async def add_user(
+    await collection.create_index('id', unique=True)
-    discord_id: int=0,
+    await collection.create_index('discord_id', unique=True)
-    tags: list=None,
+    await collection.create_index('api_key', unique=True)
-) -> dict:
+
-    """Adds a new key to the database"""
+async def add_user(discord_id: int = 0, tags: list = None) -> dict:
+    """Adds a new user to the MongoDB collection."""
 
     chars = string.ascii_letters + string.digits
 
@@ -53,63 +45,73 @@ async def add_user(
         'id': str(uuid.uuid4()),
         'api_key': key,
         'created_at': int(time.time()),
-        'last_used': 0,
+        'ban_reason': '',
-        'uses_count': 0,
-        'tokens_generated': 0,
         'active': True,
         'discord_id': discord_id,
         'credit': 0,
-        'tags': '/'.join(tags)
+        'tags': '/'.join(tags),
+        'usage': {
+            'events': [],
+            'num_tokens': 0
+        }
     }
 
-    await databases.insert_dict(new_user, 'users', 'users')
+    collection = get_mongo('users')
+    await collection.insert_one(new_user)
+
     return new_user
 
-async def get_user(
+async def get_user(by_id: str = '', by_discord_id: int = 0, by_api_key: str = ''):
-    by_id: str='',
+    """Retrieve a user from the MongoDB collection."""
-    by_discord_id: int=0,
-):
-    users_db = await databases.connect('users')
 
-    async with users_db.execute(
+    collection = get_mongo('users')
-        'SELECT * FROM users WHERE id = :id OR discord_id = :discord_id',
+    query = {
-        {'id': by_id, 'discord_id': by_discord_id}
+        '$or': [
-    ) as cursor:
+            {'id': by_id},
-        async for row in cursor:
+            {'discord_id': by_discord_id},
-            result = await databases.row_to_dict(row, cursor)
+            {'api_key': by_api_key},
-            return result
+        ]
-
+    }
-    return None
+    return await collection.find_one(query)
 
 async def get_all_users():
-    users_db = await databases.connect('users')
+    """Retrieve all users from the MongoDB collection."""
-    results = []
 
-    async with users_db.execute(
+    collection = get_mongo('users')
-        'SELECT * FROM users'
+    return list(await collection.find())
-    ) as cursor:
-        async for row in cursor:
-            result = await databases.row_to_dict(row, cursor)
-            results.append(result)
 
-    return results
+async def user_used_api(user_id: str, num_tokens: int = 0, model='', ip_address: str = '', user_agent: str = '') -> None:
+    """Update the stats of a user."""
+
+    collection = get_mongo('users')
+    user = await get_user(by_id=user_id)
+
+    if not user:
+        raise ValueError('User not found.')
+
+    usage = user['usage']
+    usage['events'].append({
+        'timestamp': time.time(),
+        'ip_address': ip_address,
+        'user_agent': user_agent,
+        'model': model,
+        'num_tokens': num_tokens
+    })
+
+    usage['num_tokens'] += num_tokens
+
+    await collection.update_one({'id': user_id}, {'$set': {'usage': usage}})
 
 async def demo():
     await prepare()
 
-    users = await get_all_users()
-    print(users)
-
     example_id = 133769420
-    user = await add_user(discord_id=example_id)
-    print(user)
-
-    del user
-    print('Fetching user...')
-
     user = await get_user(by_discord_id=example_id)
-    print(user['api_key'])
+    print(user)
+    uid = await user['id']
+
+    await user_used_api(uid, model='gpt-5', num_tokens=42, ip_address='9.9.9.9', user_agent='Mozilla/5.0')
+    # print(user)
 
 if __name__ == '__main__':
     asyncio.run(demo())
-    os.system(f'pkill -f {os.path.basename(__file__)}')

tests/__main__.py

@@ -53,7 +53,7 @@ def test_api(model: str=MODEL, messages: List[dict]=None) -> dict:
     )
     response.raise_for_status()
 
-    return response
+    return response.text
 
 def test_library():
     """Tests if the api_endpoint is working with the Python library."""
@@ -77,5 +77,5 @@ def test_all():
     # print(test_library())
 
 if __name__ == '__main__':
-    api_endpoint = 'https://api.nova-oss.com'
+    # api_endpoint = 'https://api.nova-oss.com'
     test_all()