diff --git a/.gitignore b/.gitignore index 808c342..af85215 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ docker-compose.yml venv/ -.env \ No newline at end of file +.env +*.db +__pycache__/ \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index 68b8c81..f5ca9a0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,7 +2,7 @@ FROM python:3-bullseye RUN apt update && apt install -y ffmpeg WORKDIR /app -COPY ./main.py /app +COPY ./*.py /app COPY ./requirements.txt /app RUN pip install -r requirements.txt CMD [ "python3", "/app/main.py" ] \ No newline at end of file diff --git a/database.py b/database.py new file mode 100644 index 0000000..045c615 --- /dev/null +++ b/database.py @@ -0,0 +1,113 @@ +import sqlite3 +import json + +def init_database(): + conn = sqlite3.connect("users.db") + c = conn.cursor() + c.execute(""" + CREATE TABLE IF NOT EXISTS users ( + chat_id TEXT PRIMARY KEY, + context TEXT, + usage_chatgpt INTEGER, + usage_whisper INTEGER, + usage_dalle INTEGER, + whisper_to_chat INTEGER, + assistant_voice_chat INTEGER, + temperature REAL, + max_context INTEGER + ) + """) + conn.commit() + conn.close() + +def get_user(chat_id: str): + conn = sqlite3.connect("users.db") + c = conn.cursor() + c.execute("SELECT * FROM users WHERE chat_id = ?", (chat_id,)) + user = c.fetchone() + conn.close() + if user: + return { + "context": json.loads(user[1]), + "usage": { + "chatgpt": user[2], + "whisper": user[3], + "dalle": user[4] + }, + "options": { + "whisper_to_chat": bool(user[5]), + "assistant_voice_chat": bool(user[6]), + "temperature": user[7], + "max-context": user[8] + } + } + return None + +def add_user(chat_id: str, user_data): + conn = sqlite3.connect("users.db") + c = conn.cursor() + c.execute(""" + INSERT INTO users ( + chat_id, context, usage_chatgpt, usage_whisper, usage_dalle, + whisper_to_chat, assistant_voice_chat, temperature, max_context + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, ( + chat_id, + json.dumps(user_data["context"]), + user_data["usage"]["chatgpt"], + user_data["usage"]["whisper"], + user_data["usage"]["dalle"], + int(user_data["options"]["whisper_to_chat"]), + int(user_data["options"]["assistant_voice_chat"]), + user_data["options"]["temperature"], + user_data["options"]["max-context"] + )) + conn.commit() + conn.close() + +def update_user(chat_id: str, user_data): + conn = sqlite3.connect("users.db") + c = conn.cursor() + c.execute(""" + UPDATE users + SET + context = ?, + usage_chatgpt = ?, + usage_whisper = ?, + usage_dalle = ?, + whisper_to_chat = ?, + assistant_voice_chat = ?, + temperature = ?, + max_context = ? + WHERE chat_id = ? + """, ( + json.dumps(user_data["context"]), + user_data["usage"]["chatgpt"], + user_data["usage"]["whisper"], + user_data["usage"]["dalle"], + int(user_data["options"]["whisper_to_chat"]), + int(user_data["options"]["assistant_voice_chat"]), + user_data["options"]["temperature"], + user_data["options"]["max-context"], + chat_id + )) + conn.commit() + conn.close() + +def get_total_usage(): + conn = sqlite3.connect("users.db") + c = conn.cursor() + c.execute(""" + SELECT + SUM(usage_chatgpt) AS total_chatgpt, + SUM(usage_whisper) AS total_whisper, + SUM(usage_dalle) AS total_dalle + FROM users + """) + total_usage = c.fetchone() + conn.close() + return { + "chatgpt": total_usage[0], + "whisper": total_usage[1], + "dalle": total_usage[2] + } \ No newline at end of file diff --git a/main.py b/main.py index 9795ad6..b4218d7 100644 --- a/main.py +++ b/main.py @@ -1,89 +1,91 @@ import os -import re import openai import logging -import asyncio -import math +import database from dotenv import load_dotenv from pydub import AudioSegment from telegram import Update from functools import wraps from telegram.constants import ChatAction from functools import wraps -from telegram.error import BadRequest, RetryAfter, TimedOut from telegram.ext import ApplicationBuilder, CommandHandler, ContextTypes, MessageHandler, filters, CallbackQueryHandler from telegram import InlineKeyboardButton, InlineKeyboardMarkup +language_models = { + "en": "tts_models/multilingual/multi-dataset/your_tts", + "fr": "tts_models/multilingual/multi-dataset/your_tts", + "pt": "tts_models/multilingual/multi-dataset/your_tts", + "pt-br": "tts_models/multilingual/multi-dataset/your_tts", + "es": "tts_models/es/css10/vits", +} + logging.basicConfig( format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO ) logger = logging.getLogger(__name__) +# Envrionment Variables Load load_dotenv() - if os.environ.get("OPENAI_API_KEY") is None: print("OpenAI_API_KEY is not set in.env file or OPENAI_API_KEY environment variable is not set") exit(1) - ALLOWED_USERS=os.environ.get("BOT_ALLOWED_USERS").split(",") SYSTEM_PROMPT=os.environ.get("CHATGPT_SYSTEM_PROMPT") TEMPERATURE=os.environ.get("CHATGPT_TEMPERATURE") MODEL=os.environ.get("OPENAI_MODEL") WHISPER_TO_CHAT=bool(int(os.environ.get("WHISPER_TO_CHAT"))) MAX_USER_CONTEXT=int(os.environ.get("CHATGPT_MAX_USER_CONTEXT")) - openai.api_key = os.environ.get("OPENAI_API_KEY") -users = { - "userid": { - "context": [], - "usage": { - "chatgpt": 0, - "whisper": 0, - "dalle": 0, - }, - "options": { - "whisper-to-chat": WHISPER_TO_CHAT, - "temperature": 0.9, - "max-context": 5 + +async def getUserData(chat_id): + # Initialize user if not present + user_data = database.get_user(chat_id) + if not user_data: + user_data = { + "context": [], + "usage": {"chatgpt": 0, "whisper": 0, "dalle": 0}, + "options": { + "whisper_to_chat": WHISPER_TO_CHAT, + "assistant_voice_chat": False, + "temperature": float(TEMPERATURE), + "max-context": MAX_USER_CONTEXT + } } - }, -} + database.add_user(chat_id, user_data) + user_data = database.get_user(chat_id) + return user_data def restricted(func): @wraps(func) async def wrapped(update, context, *args, **kwargs): - user_id = update.effective_user.id - if str(user_id) not in ALLOWED_USERS: + if str(update.effective_user.id) not in ALLOWED_USERS: if "*" != ALLOWED_USERS[0]: - print(f"Unauthorized access denied for {user_id}.") + print(f"Unauthorized access denied for {update.effective_user.id}.") return else: - if not f"{update.effective_chat.id}" in users: - users[f"{update.effective_chat.id}"] = {"context": [], "usage": {"chatgpt": 0,"whisper": 0,"dalle": 0,}, "options": {"whisper-to-chat": WHISPER_TO_CHAT, "temperature": float(TEMPERATURE), "max-context": MAX_USER_CONTEXT}} + _ = await getUserData(update.effective_chat.id) return await func(update, context, *args, **kwargs) return wrapped -async def messageGPT(text: str, chat_id: str): - # Initialize user if not present - if chat_id not in users: - users[chat_id] = {"context": [], "usage": {"chatgpt": 0,"whisper": 0,"dalle": 0,}, "options": {"whisper-to-chat": WHISPER_TO_CHAT, "temperature": float(TEMPERATURE), "max-context": MAX_USER_CONTEXT}} +async def messageGPT(text: str, chat_id: str, user_name="User"): + user_data = await getUserData(chat_id) # Update context - user_context = users[chat_id]["context"] - user_context.append({"role": "user", "content": text}) - if len(user_context) > users[chat_id]["options"]["max-context"]: - user_context.pop(0) + user_data['context'].append({"role": "user", "content": text}) + if len(user_data['context']) > user_data["options"]["max-context"]: + user_data['context'].pop(0) # Interact with ChatGPT API and stream the response response = None try: response = openai.ChatCompletion.create( model=MODEL, - messages=[{"role": "system", "content": SYSTEM_PROMPT}] + user_context, - temperature=users[chat_id]["options"]["temperature"], + messages=[{"role": "system", "content": f"You are chatting with {user_name}. {SYSTEM_PROMPT}"}] + user_data['context'], + temperature=user_data["options"]["temperature"], ) - except: + except Exception as e: + print(e) return "There was a problem with OpenAI, so I can't answer you." # Initialize variables for streaming @@ -94,26 +96,29 @@ async def messageGPT(text: str, chat_id: str): assistant_message = "There was a problem with OpenAI. Maybe your prompt is forbidden? They like to censor a lot!" # Update context - user_context.append({"role": "assistant", "content": assistant_message}) - if len(user_context) > users[chat_id]["options"]["max-context"]: - user_context.pop(0) + user_data['context'].append({"role": "assistant", "content": assistant_message}) + if len(user_data['context']) > user_data["options"]["max-context"]: + user_data['context'].pop(0) # Update usage - users[chat_id]["usage"]['chatgpt'] += int(response['usage']['total_tokens']) - + user_data["usage"]['chatgpt'] += int(response['usage']['total_tokens']) + + # Update the user data in the database + database.update_user(chat_id, user_data) return assistant_message @restricted -async def start(update: Update, context: ContextTypes.DEFAULT_TYPE): - if not f"{update.effective_chat.id}" in users: - users[f"{update.effective_chat.id}"] = {"context": [], "usage": {"chatgpt": 0,"whisper": 0,"dalle": 0,}, "options": {"whisper-to-chat": WHISPER_TO_CHAT, "temperature": float(TEMPERATURE), "max-context": MAX_USER_CONTEXT}} - await context.bot.send_message(chat_id=update.effective_chat.id, text="I'm a bot, please talk to me!") +async def start(update: Update, context: ContextTypes.DEFAULT_TYPE): + _ = await getUserData(update.effective_chat.id) + await context.bot.send_message(chat_id=update.effective_chat.id, text="Hello, how can I assist you today?") @restricted async def imagine(update: Update, context: ContextTypes.DEFAULT_TYPE): - users[f"{update.effective_chat.id}"]["usage"]['dalle'] += 1 - await context.bot.send_chat_action(chat_id=update.effective_chat.id, action=ChatAction.TYPING) + user_data = await getUserData(update.effective_chat.id) + user_data["usage"]['dalle'] += 1 + database.update_user(update.effective_chat.id, user_data) + await context.bot.send_chat_action(chat_id=update.effective_chat.id, action=ChatAction.TYPING) response = openai.Image.create( prompt=update.message.text, n=1, @@ -122,7 +127,8 @@ async def imagine(update: Update, context: ContextTypes.DEFAULT_TYPE): try: image_url = response['data'][0]['url'] await context.bot.send_message(chat_id=update.effective_chat.id, text=image_url) - except: + except Exception as e: + print(e) await context.bot.send_message(chat_id=update.effective_chat.id, text="Error generating. Your prompt may contain text that is not allowed by OpenAI safety system.") @restricted @@ -130,24 +136,28 @@ async def attachment(update: Update, context: ContextTypes.DEFAULT_TYPE): # Initialize variables chat_id = update.effective_chat.id await context.bot.send_chat_action(chat_id=chat_id, action=ChatAction.TYPING) - users[f"{chat_id}"]["usage"]['whisper'] = 0 + + # Get user data or initialize if not present + user_data = await getUserData(chat_id) + + #users[f"{chat_id}"]["usage"]['whisper'] = 0 transcript = {'text': ''} audioMessage = False # Check if the attachment is a voice message if update.message.voice: - users[f"{chat_id}"]["usage"]['whisper'] += update.message.voice.duration + user_data["usage"]['whisper'] += update.message.voice.duration file_id = update.message.voice.file_id file_format = "ogg" audioMessage = True # Check if the attachment is a video elif update.message.video: - users[f"{chat_id}"]["usage"]['whisper'] += update.message.video.duration + user_data["usage"]['whisper'] += update.message.video.duration file_id = update.message.video.file_id file_format = "mp4" # Check if the attachment is an audio file elif update.message.audio: - users[f"{chat_id}"]["usage"]['whisper'] += update.message.audio.duration + user_data["usage"]['whisper'] += update.message.audio.duration file_id = update.message.audio.file_id file_format = "mp3" else: @@ -169,7 +179,8 @@ async def attachment(update: Update, context: ContextTypes.DEFAULT_TYPE): with open(f"{user_id}.{file_format}", "rb") as audio_file: try: transcript = openai.Audio.transcribe("whisper-1", audio_file) - except: + except Exception as e: + print(e) await context.bot.send_message(chat_id=chat_id, text="Transcript failed.") os.remove(f"{user_id}.{file_format}") return @@ -180,8 +191,8 @@ async def attachment(update: Update, context: ContextTypes.DEFAULT_TYPE): if transcript['text'] == "": transcript['text'] = "[Silence]" - if audioMessage and users[f"{chat_id}"]["options"]["whisper-to-chat"]: - chatGPT_response = await messageGPT(transcript['text'], str(chat_id)) + if audioMessage and user_data["options"]["whisper_to_chat"]: + chatGPT_response = await messageGPT(transcript['text'], str(chat_id), update.effective_user.name) transcript['text'] = "> " + transcript['text'] + "\n\n" + chatGPT_response # Check if the transcript length is longer than 4095 characters @@ -202,16 +213,15 @@ async def attachment(update: Update, context: ContextTypes.DEFAULT_TYPE): await context.bot.send_message(chat_id=chat_id, text=current_message) else: await context.bot.send_message(chat_id=chat_id, text=transcript['text']) + + # Update user data in the database + database.update_user(str(chat_id), user_data) @restricted async def chat(update: Update, context: ContextTypes.DEFAULT_TYPE): chat_id = str(update.effective_chat.id) await context.bot.send_chat_action(chat_id=chat_id, action=ChatAction.TYPING) - # Initialize user if not present - if chat_id not in users: - users[chat_id] = {"context": [], "usage": {"chatgpt": 0, "whisper": 0, "dalle": 0}} - # Check if replying and add context if hasattr(update.message.reply_to_message, "text"): user_prompt = f"In reply to: '{update.message.reply_to_message.text}' \n---\n{update.message.text}" @@ -219,24 +229,44 @@ async def chat(update: Update, context: ContextTypes.DEFAULT_TYPE): user_prompt = update.message.text # Use messageGPT function to get the response - assistant_message = await messageGPT(user_prompt, chat_id) + assistant_message = await messageGPT(user_prompt, chat_id, update.effective_user.name) await context.bot.send_message(chat_id=update.effective_chat.id, text=assistant_message) @restricted async def clear(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - if f"{update.effective_chat.id}" in users: - users[f"{update.effective_chat.id}"]["context"] = [] - print(f"Cleared context for {update.effective_chat.id}") - await update.message.reply_text(f'Your message context history was cleared.') + user_data = await getUserData(update.effective_chat.id) + if user_data: + user_data["context"] = [] + database.update_user(str(update.effective_chat.id), user_data) + print(f"Cleared context for {update.effective_user.name}") + await update.message.reply_text('Your message context history was cleared.') @restricted async def usage(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - user_info=users[f"{update.effective_chat.id}"]["usage"] - total_spent=0.0 - total_spent+=(user_info['chatgpt']/750)*0.002 - total_spent+=float(user_info['dalle'])*0.02 - total_spent+=(user_info['whisper']/60.0)*0.006 - info_message=f"""User: {update.effective_user.name}\n- Used ~{user_info["chatgpt"]} tokens with ChatGPT.\n- Generated {user_info["dalle"]} images with DALL-E.\n- Transcribed {round(float(user_info["whisper"])/60.0, 2)}min with Whisper.\n\nTotal spent: ${str(total_spent)}""" + chat_id = str(update.effective_chat.id) + user_data = database.get_user(chat_id) + user_usage = user_data["usage"] + total_usage = database.get_total_usage() + + user_spent = round((((user_usage['chatgpt'] / 750) * 0.002) + (float(user_usage['dalle']) * 0.02) + ((user_usage['whisper'] / 60.0) * 0.006)), 4) + total_spent = round((((total_usage['chatgpt'] / 750) * 0.002) + (float(total_usage['dalle']) * 0.02) + ((total_usage['whisper'] / 60.0) * 0.006)), 4) + + user_percentage = (user_spent / total_spent) * 100 if total_spent > 0 else 0 + + info_message = f"""User: {update.effective_user.name} +- Used ~{user_usage["chatgpt"]} tokens with ChatGPT. +- Generated {user_usage["dalle"]} images with DALL-E. +- Transcribed {round(float(user_usage["whisper"]) / 60.0, 2)}min with Whisper. + +Total spent: ${user_spent} ({user_percentage:.2f}% of total) + +Total usage: +- ChatGPT tokens: {total_usage["chatgpt"]} +- DALL-E images: {total_usage["dalle"]} +- Whisper transcription: {round(float(total_usage["whisper"]) / 60.0, 2)}min + +Total spent: ${total_spent}""" + await context.bot.send_message(chat_id=update.effective_chat.id, text=info_message) @restricted @@ -260,6 +290,10 @@ def generate_settings_markup(chat_id: str) -> InlineKeyboardMarkup: InlineKeyboardButton("Enable Whisper to Chat", callback_data=f"setting_enable_whisper_{chat_id}"), InlineKeyboardButton("Disable Whisper to Chat", callback_data=f"setting_disable_whisper_{chat_id}") ], + [ + InlineKeyboardButton("Enable assistant voice", callback_data=f"setting_enable_voice_{chat_id}"), + InlineKeyboardButton("Disable assistant voice", callback_data=f"setting_disable_voice_{chat_id}") + ], [ InlineKeyboardButton("Increase Context", callback_data=f"setting_increase_context_{chat_id}"), InlineKeyboardButton("Decrease Context", callback_data=f"setting_decrease_context_{chat_id}") @@ -274,22 +308,37 @@ async def settings(update: Update, context: ContextTypes.DEFAULT_TYPE): await context.bot.send_message(chat_id=chat_id, text="Settings:", reply_markup=settings_markup) async def settings_callback(update: Update, context: ContextTypes.DEFAULT_TYPE): + user_data = await getUserData(update.effective_chat.id) query = update.callback_query action, chat_id = query.data.rsplit("_", 1) + + # Temperature if action.startswith("setting_increase_temperature"): - users[chat_id]["options"]["temperature"] = min(users[chat_id]["options"]["temperature"] + 0.1, 1) + user_data["options"]["temperature"] = min(user_data["options"]["temperature"] + 0.1, 1) elif action.startswith("setting_decrease_temperature"): - users[chat_id]["options"]["temperature"] = max(users[chat_id]["options"]["temperature"] - 0.1, 0) + user_data["options"]["temperature"] = max(user_data["options"]["temperature"] - 0.1, 0) + + # Whisper to GPT elif action.startswith("setting_enable_whisper"): print(f"enabling whisper for {chat_id}") - users[chat_id]["options"]["whisper-to-chat"] = True + user_data["options"]["whisper_to_chat"] = True elif action.startswith("setting_disable_whisper"): print(f"disabling whisper for {chat_id}") - users[chat_id]["options"]["whisper-to-chat"] = False + user_data["options"]["whisper_to_chat"] = False + + # TTS + elif action.startswith("setting_enable_voice"): + print(f"enabling voice for {chat_id}") + user_data["options"]["assistant_voice_chat"] = True + elif action.startswith("setting_disable_voice"): + print(f"disabling voice for {chat_id}") + user_data["options"]["assistant_voice_chat"] = False + + # Context elif action.startswith("setting_increase_context"): - users[chat_id]["options"]["max-context"] = min(users[chat_id]["options"]["max-context"] + 1, MAX_USER_CONTEXT) + user_data["options"]["max-context"] = min(user_data["options"]["max-context"] + 1, MAX_USER_CONTEXT) elif action.startswith("setting_decrease_context"): - users[chat_id]["options"]["max-context"] = max(users[chat_id]["options"]["max-context"] - 1, 1) + user_data["options"]["max-context"] = max(user_data["options"]["max-context"] - 1, 1) settings_markup = generate_settings_markup(chat_id) await query.edit_message_text(text="Choose a setting option:", reply_markup=settings_markup) @@ -298,13 +347,15 @@ async def settings_callback(update: Update, context: ContextTypes.DEFAULT_TYPE): await context.bot.delete_message(chat_id=query.message.chat_id, message_id=query.message.message_id) # Send a message displaying the updated settings - settings_message = f"""Updated settings:\n\nTemperature: {users[chat_id]['options']['temperature']}\nWhisper to Chat: {users[chat_id]['options']['whisper-to-chat']}\nContext Length: {users[chat_id]["options"]["max-context"]}""" + settings_message = f"""Updated settings:\n\nTemperature: {user_data['options']['temperature']}\nWhisper to Chat: {user_data['options']['whisper_to_chat']}\nAssistant voice: {user_data['options']['assistant_voice_chat']}\nContext Length: {user_data["options"]["max-context"]}""" await context.bot.send_message(chat_id=chat_id, text=settings_message) if __name__ == '__main__': + database.init_database() + try: ALLOWED_USERS=os.environ.get("BOT_ALLOWED_USERS").split(",") - except: + except (Exception): ALLOWED_USERS=ALLOWED_USERS print(f"Allowed users: {ALLOWED_USERS}") print(f"System prompt: {SYSTEM_PROMPT}") diff --git a/requirements.txt b/requirements.txt index ce8352c..b056ea8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ -openai -python-telegram-bot -pydub -python-dotenv -asyncio \ No newline at end of file +openai==0.27.2 +pydub==0.25.1 +python-dotenv==1.0.0 +python-telegram-bot==20.2