Merge pull request #3 from ijwfly/feature/allowed-list

Access control

README.md

@@ -13,6 +13,7 @@ This GitHub repository contains the implementation of a telegram bot, designed t
 7. **Model Support**: The bot supports both gpt-3.5-turbo and gpt-4 models with the capability to switch between them on-the-fly.
 8. **Customizable System Prompts**: Enables the user to initiate conversations with custom system prompts to shape the bot's behavior.
 9. **Context Window Size Customization**: The bot provides a feature to customize the maximum context window size. This allows users to set the context size for gpt-3.5-turbo and gpt-4 models individually, enabling more granular control over usage costs. This feature is particularly useful for managing API usage and optimizing the balance between cost and performance.
+10. **Access Control**: The bot includes a feature for access control. Each user is assigned a role (stranger, basic, advanced, admin), and depending on the role, they gain access to the bot. Role management is carried out through a messaging mechanism, with inline buttons sent to the admin for role changes.
 
 The purpose of this telegram bot is to create a ChatGpt-like user-friendly platform for interacting with GPT models. The repository is open for exploration, feedback, and contributions.
 

app/bot/settings_menu.py

@@ -3,7 +3,7 @@
 from aiogram import Bot, types, Dispatcher
 
 from app.storage.db import User, DB
-
+from app.storage.user_role import check_access_conditions
 
 GPT_MODELS_OPTIONS = {
     'gpt-3.5-turbo': 'GPT-3.5',
@@ -12,6 +12,10 @@
 }
 
 
+SETTINGS_PREFIX = 'settings'
+HIDE_COMMAND = 'hide'
+
+
 class VisibleOptionsSetting:
     def __init__(self, model_field: str, options):
         self.model_field = model_field
@@ -90,35 +94,50 @@ def __init__(self, bot: Bot, dispatcher: Dispatcher, db: DB):
             'use_functions': OnOffSetting('Use functions', 'use_functions'),
             'auto_summarize': OnOffSetting('Auto summarize', 'auto_summarize'),
         }
-        self.dispatcher.register_callback_query_handler(self.process_callback, lambda c: c.data in self.settings or c.data == 'settings.hide')
+        self.minimum_required_roles = {
+            'current_model': settings.USER_ROLE_CHOOSE_MODEL,
+        }
+        self.dispatcher.register_callback_query_handler(self.process_callback, lambda c: SETTINGS_PREFIX in c.data)
 
-    async def send_settings(self, message: types.Message):
-        user = await self.db.get_or_create_user(message.from_user.id)
+    async def send_settings(self, message: types.Message, user: User):
         await message.answer("Settings:", reply_markup=self.get_keyboard(user), parse_mode=types.ParseMode.MARKDOWN)
 
+    def is_setting_available_for_user(self, setting_name: str, user: User):
+        mininum_required_role = self.minimum_required_roles.get(setting_name)
+        if mininum_required_role and not check_access_conditions(mininum_required_role, user.role):
+            return False
+        return True
+
     def get_keyboard(self, user: User):
         keyboard = types.InlineKeyboardMarkup()
         for setting_name, setting_obj in self.settings.items():
+            if not self.is_setting_available_for_user(setting_name, user):
+                continue
+
             text = setting_obj.get_button_string(user)
-            keyboard.add(types.InlineKeyboardButton(text=text, callback_data=setting_name))
-        keyboard.add(types.InlineKeyboardButton(text='Hide settings', callback_data='settings.hide'))
+            keyboard.add(types.InlineKeyboardButton(text=text, callback_data=f'{SETTINGS_PREFIX}.{setting_name}'))
+        keyboard.add(types.InlineKeyboardButton(text='Hide settings', callback_data=f'{SETTINGS_PREFIX}.{HIDE_COMMAND}'))
         return keyboard
 
-    def toggle_setting(self, user: User, setting: str):
-        setting = self.settings[setting]
-        user = setting.toggle(user)
+    def toggle_setting(self, user: User, setting_name: str):
+        if not self.is_setting_available_for_user(setting_name, user):
+            return user
+        setting_name = self.settings[setting_name]
+        user = setting_name.toggle(user)
         return user
 
     async def process_callback(self, callback_query: types.CallbackQuery):
-        if callback_query.data == 'settings.hide':
+        _, command = callback_query.data.split('.')
+        if command == HIDE_COMMAND:
             await self.bot.delete_message(
                 chat_id=callback_query.from_user.id,
                 message_id=callback_query.message.message_id
             )
             await self.bot.answer_callback_query(callback_query.id)
         else:
+            setting = command
             user = await self.db.get_or_create_user(callback_query.from_user.id)
-            user = self.toggle_setting(user, callback_query.data)
+            user = self.toggle_setting(user, setting)
             await self.db.update_user(user)
 
             await self.bot.answer_callback_query(callback_query.id)

app/bot/telegram_bot.py

@@ -6,12 +6,14 @@
 from app.bot.dialog_manager import DialogUtils
 from app.bot.settings_menu import Settings
 from app.bot.user_middleware import UserMiddleware
+from app.bot.user_role_manager import UserRoleManager
 from app.bot.utils import TypingWorker, detect_and_extract_code, get_username, message_is_forward, get_hide_button
 from app.context.context_manager import build_context_manager
 from app.openai_helpers.function_storage import FunctionStorage
 from app.openai_helpers.utils import calculate_completion_usage_price, calculate_whisper_usage_price
 from app.openai_helpers.whisper import get_audio_speech_to_text
 from app.storage.db import DBFactory, User
+from app.storage.user_role import check_access_conditions
 from app.openai_helpers.chatgpt import ChatGPT, GptModel
 
 from aiogram.utils.exceptions import CantParseEntities
@@ -37,13 +39,15 @@ def __init__(self, bot: Bot, dispatcher: Dispatcher, function_storage: FunctionS
 
         # initialized in on_startup
         self.settings = None
+        self.role_manager = None
 
     async def on_startup(self, _):
         self.db = await DBFactory().create_database(
             settings.POSTGRES_USER, settings.POSTGRES_PASSWORD,
             settings.POSTGRES_HOST, settings.POSTGRES_PORT, settings.POSTGRES_DATABASE
         )
         self.settings = Settings(self.bot, self.dispatcher, self.db)
+        self.role_manager = UserRoleManager(self.bot, self.dispatcher, self.db)
         self.dispatcher.middleware.setup(UserMiddleware(self.db))
 
         commands = [
@@ -185,6 +189,10 @@ async def reset_dialog(self, message: types.Message, user: User):
         await message.answer('👌')
 
     async def set_current_model(self, message: types.Message, user: User):
+        if not check_access_conditions(settings.USER_ROLE_CHOOSE_MODEL, user.role):
+            await message.answer(f'Your model is {user.current_model}. You have no permissions to change model')
+            return
+
         model = GptModel.GPT_35_TURBO if message.get_command() == '/gpt3' else GptModel.GPT_4
         user.current_model = model
         await self.db.update_user(user)
@@ -209,9 +217,9 @@ async def get_usage(self, message: types.Message, user: User):
             message, '\n'.join(result), types.ParseMode.MARKDOWN, reply_markup=get_hide_button()
         )
 
-    async def open_settings(self, message: types.Message):
+    async def open_settings(self, message: types.Message, user: User):
         await self.bot.delete_message(
             chat_id=message.from_user.id,
             message_id=message.message_id,
         )
-        await self.settings.send_settings(message)
+        await self.settings.send_settings(message, user)

app/bot/user_middleware.py

@@ -1,16 +1,46 @@
+import settings
+from app.bot.user_role_manager import UserRoleManager
+from app.storage.db import DB
+from app.storage.user_role import check_access_conditions
+
 from aiogram import types
+from aiogram.dispatcher.handler import CancelHandler
 from aiogram.dispatcher.middlewares import BaseMiddleware
 
-from app.storage.db import DB
-
 
 class UserMiddleware(BaseMiddleware):
     def __init__(self, db: DB):
         super().__init__()
         self.db = db
 
     async def on_pre_process_message(self, message: types.Message, data: dict):
+        is_new_user = False
+
         user_id = message.from_user.id
-        # Здесь вы можете получить пользователя из базы данных
-        user = await self.db.get_or_create_user(user_id)
+        user = await self.db.get_user(user_id)
+        if user is None:
+            user = await self.db.create_user(user_id, settings.USER_ROLE_DEFAULT)
+            is_new_user = True
+
+        if user.role is None:
+            user.role = settings.USER_ROLE_DEFAULT
+            await self.db.update_user(user)
+
+        full_name = message.from_user.full_name
+        username = message.from_user.username
+        if user.full_name != full_name or user.username != username:
+            user.full_name = full_name
+            user.username = username
+            await self.db.update_user(user)
+
+        if settings.ENABLE_USER_ROLE_MANAGER_CHAT and is_new_user:
+            await UserRoleManager.send_new_user_to_admin(message, user)
+
+        user_have_access = check_access_conditions(settings.USER_ROLE_BOT_ACCESS, user.role)
+        if not user_have_access:
+            await message.answer(
+                "You currently don't have access to this bot. You will be notified once the admin grants you access."
+            )
+            raise CancelHandler()
+
         data['user'] = user

app/bot/user_role_manager.py

@@ -0,0 +1,76 @@
+from aiogram import types, Bot, Dispatcher
+
+import settings
+from app.storage.db import User, DB
+from app.storage.user_role import UserRole, check_access_conditions
+
+
+SET_ROLE_COMMAND = 'setrole'
+UPDATE_INFO_COMMAND = 'updinfo'
+
+
+class UserRoleManager:
+    def __init__(self, bot: Bot, dispatcher: Dispatcher, db: DB):
+        self.bot = bot
+        self.dispatcher = dispatcher
+        self.db = db
+        self.dispatcher.register_callback_query_handler(
+            self.setrole_callback, lambda c: SET_ROLE_COMMAND in c.data,
+        )
+        self.dispatcher.register_callback_query_handler(
+            self.updaterole_callback, lambda c: UPDATE_INFO_COMMAND in c.data,
+        )
+
+    @staticmethod
+    def get_keyboard(user: User):
+        keyboard = types.InlineKeyboardMarkup()
+
+        for role in UserRole:
+            callback_data = f'{SET_ROLE_COMMAND}.{user.telegram_id}.{role.value}'
+            if role == user.role:
+                keyboard.add(types.InlineKeyboardButton(text=f'<{role.value}>', callback_data=callback_data))
+            else:
+                keyboard.add(types.InlineKeyboardButton(text=role.value, callback_data=callback_data))
+        keyboard.add(types.InlineKeyboardButton(text='🔄', callback_data=f'{UPDATE_INFO_COMMAND}.{user.telegram_id}'))
+        return keyboard
+
+    @staticmethod
+    def user_to_string(user):
+        result = [f'*User Id*: {user.id}', f'*Telegram Id*: {user.telegram_id}']
+        if user.full_name:
+            result.append(f'*Full name*: {user.full_name}')
+        if user.username:
+            result.append(f'*Username*: @{user.username}')
+        result.append(f'*Role*: {user.role.value}')
+        return '\n'.join(result)
+
+    @classmethod
+    async def send_new_user_to_admin(cls, message: types.Message, user: User):
+        bot = message.bot
+        text = cls.user_to_string(user)
+        await bot.send_message(
+            settings.USER_ROLE_MANAGER_CHAT_ID, text, reply_markup=cls.get_keyboard(user), parse_mode=types.ParseMode.MARKDOWN
+        )
+
+    async def update_message(self, message: types.Message, user: User):
+        text = self.user_to_string(user)
+        await message.edit_text(text, reply_markup=self.get_keyboard(user), parse_mode=types.ParseMode.MARKDOWN)
+
+    async def setrole_callback(self, callback_query: types.CallbackQuery):
+        command, tg_user_id, role_value = callback_query.data.split('.')
+        tg_user_id = int(tg_user_id)
+        user = await self.db.get_user(tg_user_id)
+        user_had_access = check_access_conditions(settings.USER_ROLE_BOT_ACCESS, user.role)
+        user.role = UserRole(role_value)
+        await self.db.update_user(user)
+        await self.bot.answer_callback_query(callback_query.id)
+        await self.update_message(callback_query.message, user)
+        if check_access_conditions(settings.USER_ROLE_BOT_ACCESS, user.role) and not user_had_access:
+            await self.bot.send_message(tg_user_id, f'You have been granted access to the bot.')
+
+    async def updaterole_callback(self, callback_query: types.CallbackQuery):
+        command, tg_user_id = callback_query.data.split('.')
+        tg_user_id = int(tg_user_id)
+        user = await self.db.get_user(tg_user_id)
+        await self.bot.answer_callback_query(callback_query.id)
+        await self.update_message(callback_query.message, user)

app/openai_helpers/utils.py

@@ -2,7 +2,6 @@
 import openai
 
 
-
 COMPLETION_PRICE = {
     'gpt-3.5-turbo': (Decimal('0.0015'), Decimal('0.002')),
     'gpt-3.5-turbo-16k': (Decimal('0.003'), Decimal('0.004')),

app/storage/db.py

@@ -1,13 +1,15 @@
 import json
 from datetime import datetime
 from enum import Enum
-from typing import List
+from typing import List, Optional
+
+import settings
+from app.openai_helpers.chatgpt import DialogMessage, CompletionUsage
+from app.storage.user_role import UserRole
 
 import asyncpg
 import pydantic
 
-from app.openai_helpers.chatgpt import DialogMessage, CompletionUsage
-
 
 class User(pydantic.BaseModel):
     id: int
@@ -18,6 +20,9 @@ class User(pydantic.BaseModel):
     voice_as_prompt: bool
     use_functions: bool
     auto_summarize: bool
+    full_name: Optional[str]
+    username: Optional[str]
+    role: Optional[UserRole]
 
 
 class MessageType(Enum):
@@ -45,7 +50,7 @@ def __init__(self, connection_pool: asyncpg.Pool):
     async def get_or_create_user(self, telegram_user_id):
         user = await self.get_user(telegram_user_id)
         if user is None:
-            user = await self.create_user(telegram_user_id)
+            user = await self.create_user(telegram_user_id, settings.USER_ROLE_DEFAULT)
         return user
 
     async def get_user(self, telegram_user_id):
@@ -58,17 +63,19 @@ async def get_user(self, telegram_user_id):
     async def update_user(self, user: User):
         sql = '''UPDATE chatgpttg.user 
         SET current_model = $1, gpt_mode = $2, forward_as_prompt = $3,
-        voice_as_prompt = $4, use_functions = $5, auto_summarize = $6 WHERE id = $7 RETURNING *'''
+        voice_as_prompt = $4, use_functions = $5, auto_summarize = $6,
+        full_name = $7, username = $8, role = $9 WHERE id = $10 RETURNING *'''
         return User(**await self.connection_pool.fetchrow(
             sql, user.current_model, user.gpt_mode, user.forward_as_prompt,
-            user.voice_as_prompt, user.use_functions, user.auto_summarize, user.id,
+            user.voice_as_prompt, user.use_functions, user.auto_summarize,
+            user.full_name, user.username, user.role.value, user.id,
         ))
 
-    async def create_user(self, telegram_user_id):
-        sql = 'INSERT INTO chatgpttg.user (telegram_id) VALUES ($1) RETURNING *'
-        return User(**await self.connection_pool.fetchrow(sql, telegram_user_id))
+    async def create_user(self, telegram_user_id: int, role: UserRole):
+        sql = 'INSERT INTO chatgpttg.user (telegram_id, role) VALUES ($1, $2) RETURNING *'
+        return User(**await self.connection_pool.fetchrow(sql, telegram_user_id, role.value))
 
-    async def get_telegram_message(self, tg_chat_id, tg_message_id):
+    async def get_telegram_message(self, tg_chat_id: int, tg_message_id: int):
         sql = 'SELECT * FROM chatgpttg.message WHERE tg_chat_id = $1 AND tg_message_id = $2'
         tg_message_record = await self.connection_pool.fetchrow(sql, tg_chat_id, tg_message_id)
         if tg_message_record is None:

app/storage/user_role.py

@@ -0,0 +1,15 @@
+from enum import Enum
+
+
+class UserRole(Enum):
+    ADMIN = 'admin'
+    ADVANCED = 'advanced'
+    BASIC = 'basic'
+    STRANGER = 'stranger'
+
+
+ROLE_ORDER = [UserRole.STRANGER, UserRole.BASIC, UserRole.ADVANCED, UserRole.ADMIN]
+
+
+def check_access_conditions(required_role: UserRole, user_role: UserRole) -> bool:
+    return ROLE_ORDER.index(user_role) >= ROLE_ORDER.index(required_role)

main.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import settings
 from app.bot.queued_dispatcher import QueuedDispatcher
 from app.bot.telegram_bot import TelegramBot
@@ -11,12 +13,15 @@
 dp = QueuedDispatcher(bot)
 
 
-def setup_function_storage() -> FunctionStorage:
+def setup_function_storage() -> Optional[FunctionStorage]:
     functions = []
 
     if settings.ENABLE_WOLFRAMALPHA:
         functions.append(query_wolframalpha)
 
+    if not functions:
+        return None
+
     function_storage = FunctionStorage()
     for function in functions:
         function_storage.register(function)

migrations/sql/0003_add_user_roles.sql

@@ -0,0 +1,14 @@
+ALTER TABLE chatgpttg.user ADD COLUMN IF NOT EXISTS full_name text;
+ALTER TABLE chatgpttg.user ADD COLUMN IF NOT EXISTS username text;
+
+-- create user_role field
+DO $$
+    BEGIN
+        IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'user_roles') THEN
+            CREATE TYPE chatgpttg.user_roles AS ENUM ('admin', 'advanced', 'basic', 'stranger');
+        END IF;
+    END
+$$;
+ALTER TABLE chatgpttg.user ADD COLUMN IF NOT EXISTS role chatgpttg.user_roles;
+ALTER TABLE chatgpttg.user ADD COLUMN IF NOT EXISTS cdate timestamp WITH TIME ZONE NOT NULL default NOW();
+