update to use sqlite

.gitignore

@@ -1,3 +1,5 @@
 docker-compose.yml
 venv/
-.env
+.env
+*.db
+__pycache__/

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" ]

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]
+    }

main.py

@@ -1,89 +1,91 @@
 import os
-import re
 import openai
 import logging
-import asyncio
+import database
-import math
 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": {
+async def getUserData(chat_id):    
-        "context": [],
+    # Initialize user if not present
-        "usage": {
+    user_data = database.get_user(chat_id)
-            "chatgpt": 0,
+    if not user_data:
-            "whisper": 0,
+        user_data = {
-            "dalle": 0,
+            "context": [],
-        },
+            "usage": {"chatgpt": 0, "whisper": 0, "dalle": 0},
-        "options": {
+            "options": {
-            "whisper-to-chat": WHISPER_TO_CHAT,
+                "whisper_to_chat": WHISPER_TO_CHAT,
-            "temperature": 0.9,
+                "assistant_voice_chat": False,
-            "max-context": 5
+                "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(update.effective_user.id) not in ALLOWED_USERS:
-        if str(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:
+            _ = await getUserData(update.effective_chat.id)
-                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}}
         return await func(update, context, *args, **kwargs)
     return wrapped
 
-async def messageGPT(text: str, chat_id: str):
+async def messageGPT(text: str, chat_id: str, user_name="User"):
-    # Initialize user if not present
+    user_data = await getUserData(chat_id)
-    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}}
 
     # Update context
-    user_context = users[chat_id]["context"]
+    user_data['context'].append({"role": "user", "content": text})
-    user_context.append({"role": "user", "content": text})
+    if len(user_data['context']) > user_data["options"]["max-context"]:
-    if len(user_context) > users[chat_id]["options"]["max-context"]:
+        user_data['context'].pop(0)
-        user_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,
+            messages=[{"role": "system", "content": f"You are chatting with {user_name}. {SYSTEM_PROMPT}"}] + user_data['context'],
-            temperature=users[chat_id]["options"]["temperature"],
+            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})
+    user_data['context'].append({"role": "assistant", "content": assistant_message})
-    if len(user_context) > users[chat_id]["options"]["max-context"]:
+    if len(user_data['context']) > user_data["options"]["max-context"]:
-        user_context.pop(0)
+        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):
+async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):    
-    if not f"{update.effective_chat.id}" in users:
+    _ = await getUserData(update.effective_chat.id)
-        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="Hello, how can I assist you today?")
-    await context.bot.send_message(chat_id=update.effective_chat.id, text="I'm a bot, please talk to me!")
 
 @restricted    
 async def imagine(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    users[f"{update.effective_chat.id}"]["usage"]['dalle'] += 1
+    user_data = await getUserData(update.effective_chat.id)
-    await context.bot.send_chat_action(chat_id=update.effective_chat.id, action=ChatAction.TYPING)
+    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"]:
+    if audioMessage and user_data["options"]["whisper_to_chat"]:
-        chatGPT_response = await messageGPT(transcript['text'], str(chat_id))
+        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:
+    user_data = await getUserData(update.effective_chat.id)
-        users[f"{update.effective_chat.id}"]["context"] = []
+    if user_data:
-    print(f"Cleared context for {update.effective_chat.id}")
+        user_data["context"] = []
-    await update.message.reply_text(f'Your message context history was cleared.')
+    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"]
+    chat_id = str(update.effective_chat.id)
-    total_spent=0.0
+    user_data = database.get_user(chat_id)
-    total_spent+=(user_info['chatgpt']/750)*0.002
+    user_usage = user_data["usage"]
-    total_spent+=float(user_info['dalle'])*0.02
+    total_usage = database.get_total_usage()
-    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)}"""
+    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}")

requirements.txt

@@ -1,5 +1,4 @@
-openai
+openai==0.27.2
-python-telegram-bot
+pydub==0.25.1
-pydub
+python-dotenv==1.0.0
-python-dotenv
+python-telegram-bot==20.2
-asyncio