2023-03-06 15:20:37 +00:00
import os
import openai
import logging
2023-04-19 10:53:18 +00:00
import database
2023-03-22 11:44:08 +00:00
from dotenv import load_dotenv
2023-03-06 15:20:37 +00:00
from pydub import AudioSegment
from telegram import Update
from functools import wraps
2023-03-23 08:14:32 +00:00
from telegram . constants import ChatAction
from functools import wraps
2023-04-12 10:42:28 +00:00
from telegram . ext import ApplicationBuilder , CommandHandler , ContextTypes , MessageHandler , filters , CallbackQueryHandler
from telegram import InlineKeyboardButton , InlineKeyboardMarkup
2023-03-06 15:20:37 +00:00
2023-04-19 10:53:18 +00:00
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 " ,
}
2023-03-06 15:20:37 +00:00
logging . basicConfig (
format = ' %(asctime)s - %(name)s - %(levelname)s - %(message)s ' ,
level = logging . INFO
)
logger = logging . getLogger ( __name__ )
2023-04-19 10:53:18 +00:00
# Envrionment Variables Load
2023-03-22 11:44:08 +00:00
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 )
2023-04-12 10:42:28 +00:00
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 " ) )
2023-03-22 11:44:08 +00:00
openai . api_key = os . environ . get ( " OPENAI_API_KEY " )
2023-04-19 10:53:18 +00:00
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
}
2023-03-06 15:20:37 +00:00
}
2023-04-19 10:53:18 +00:00
database . add_user ( chat_id , user_data )
user_data = database . get_user ( chat_id )
return user_data
2023-03-06 15:20:37 +00:00
def restricted ( func ) :
@wraps ( func )
async def wrapped ( update , context , * args , * * kwargs ) :
2023-04-19 10:53:18 +00:00
if str ( update . effective_user . id ) not in ALLOWED_USERS :
2023-03-06 15:20:37 +00:00
if " * " != ALLOWED_USERS [ 0 ] :
2023-04-19 10:53:18 +00:00
print ( f " Unauthorized access denied for { update . effective_user . id } . " )
2023-03-06 15:20:37 +00:00
return
else :
2023-04-19 10:53:18 +00:00
_ = await getUserData ( update . effective_chat . id )
2023-03-06 15:20:37 +00:00
return await func ( update , context , * args , * * kwargs )
return wrapped
2023-04-19 10:53:18 +00:00
async def messageGPT ( text : str , chat_id : str , user_name = " User " ) :
user_data = await getUserData ( chat_id )
2023-04-12 10:42:28 +00:00
# Update context
2023-04-19 10:53:18 +00:00
user_data [ ' context ' ] . append ( { " role " : " user " , " content " : text } )
if len ( user_data [ ' context ' ] ) > user_data [ " options " ] [ " max-context " ] :
user_data [ ' context ' ] . pop ( 0 )
2023-04-12 10:42:28 +00:00
# Interact with ChatGPT API and stream the response
response = None
try :
response = openai . ChatCompletion . create (
model = MODEL ,
2023-04-19 10:53:18 +00:00
messages = [ { " role " : " system " , " content " : f " You are chatting with { user_name } . { SYSTEM_PROMPT } " } ] + user_data [ ' context ' ] ,
temperature = user_data [ " options " ] [ " temperature " ] ,
2023-04-12 10:42:28 +00:00
)
2023-04-19 10:53:18 +00:00
except Exception as e :
print ( e )
2023-04-12 10:42:28 +00:00
return " There was a problem with OpenAI, so I can ' t answer you. "
# Initialize variables for streaming
assistant_message = " "
if ' choices ' in response :
assistant_message = response [ ' choices ' ] [ 0 ] [ ' message ' ] [ ' content ' ]
else :
assistant_message = " There was a problem with OpenAI. Maybe your prompt is forbidden? They like to censor a lot! "
# Update context
2023-04-19 10:53:18 +00:00
user_data [ ' context ' ] . append ( { " role " : " assistant " , " content " : assistant_message } )
if len ( user_data [ ' context ' ] ) > user_data [ " options " ] [ " max-context " ] :
user_data [ ' context ' ] . pop ( 0 )
2023-04-12 10:42:28 +00:00
# Update usage
2023-04-19 10:53:18 +00:00
user_data [ " usage " ] [ ' chatgpt ' ] + = int ( response [ ' usage ' ] [ ' total_tokens ' ] )
# Update the user data in the database
database . update_user ( chat_id , user_data )
2023-04-12 10:42:28 +00:00
return assistant_message
2023-03-06 15:20:37 +00:00
@restricted
2023-04-19 10:53:18 +00:00
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? " )
2023-03-06 15:20:37 +00:00
@restricted
async def imagine ( update : Update , context : ContextTypes . DEFAULT_TYPE ) :
2023-04-19 10:53:18 +00:00
user_data = await getUserData ( update . effective_chat . id )
user_data [ " usage " ] [ ' dalle ' ] + = 1
database . update_user ( update . effective_chat . id , user_data )
2023-03-23 08:14:32 +00:00
2023-04-19 10:53:18 +00:00
await context . bot . send_chat_action ( chat_id = update . effective_chat . id , action = ChatAction . TYPING )
2023-03-06 15:20:37 +00:00
response = openai . Image . create (
prompt = update . message . text ,
n = 1 ,
size = " 1024x1024 "
)
try :
image_url = response [ ' data ' ] [ 0 ] [ ' url ' ]
await context . bot . send_message ( chat_id = update . effective_chat . id , text = image_url )
2023-04-19 10:53:18 +00:00
except Exception as e :
print ( e )
2023-03-06 15:20:37 +00:00
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
async def attachment ( update : Update , context : ContextTypes . DEFAULT_TYPE ) :
2023-03-22 11:44:08 +00:00
# Initialize variables
chat_id = update . effective_chat . id
2023-03-23 08:14:32 +00:00
await context . bot . send_chat_action ( chat_id = chat_id , action = ChatAction . TYPING )
2023-04-19 10:53:18 +00:00
# Get user data or initialize if not present
user_data = await getUserData ( chat_id )
#users[f"{chat_id}"]["usage"]['whisper'] = 0
2023-03-22 11:44:08 +00:00
transcript = { ' text ' : ' ' }
2023-04-12 10:42:28 +00:00
audioMessage = False
2023-03-22 11:44:08 +00:00
# Check if the attachment is a voice message
2023-03-07 20:06:14 +00:00
if update . message . voice :
2023-04-19 10:53:18 +00:00
user_data [ " usage " ] [ ' whisper ' ] + = update . message . voice . duration
2023-03-22 11:44:08 +00:00
file_id = update . message . voice . file_id
file_format = " ogg "
2023-04-12 10:42:28 +00:00
audioMessage = True
2023-03-22 11:44:08 +00:00
# Check if the attachment is a video
2023-03-07 20:06:14 +00:00
elif update . message . video :
2023-04-19 10:53:18 +00:00
user_data [ " usage " ] [ ' whisper ' ] + = update . message . video . duration
2023-03-22 11:44:08 +00:00
file_id = update . message . video . file_id
file_format = " mp4 "
# Check if the attachment is an audio file
elif update . message . audio :
2023-04-19 10:53:18 +00:00
user_data [ " usage " ] [ ' whisper ' ] + = update . message . audio . duration
2023-03-22 11:44:08 +00:00
file_id = update . message . audio . file_id
file_format = " mp3 "
else :
await context . bot . send_message ( chat_id = chat_id , text = " Can ' t handle such file. Reason: unknown. " )
return
# Download the file and convert it if necessary
file = await context . bot . get_file ( file_id )
2023-04-12 10:42:28 +00:00
user_id = update . effective_user . id
2023-03-22 11:44:08 +00:00
await file . download_to_drive ( f " { user_id } . { file_format } " )
if file_format == " ogg " :
ogg_audio = AudioSegment . from_file ( f " { user_id } .ogg " , format = " ogg " )
ogg_audio . export ( f " { user_id } .mp3 " , format = " mp3 " )
os . remove ( f " { user_id } .ogg " )
file_format = " mp3 "
# Transcribe the audio
with open ( f " { user_id } . { file_format } " , " rb " ) as audio_file :
2023-03-07 20:06:14 +00:00
try :
2023-03-22 11:44:08 +00:00
transcript = openai . Audio . transcribe ( " whisper-1 " , audio_file )
2023-04-19 10:53:18 +00:00
except Exception as e :
print ( e )
2023-03-22 11:44:08 +00:00
await context . bot . send_message ( chat_id = chat_id , text = " Transcript failed. " )
os . remove ( f " { user_id } . { file_format } " )
return
os . remove ( f " { user_id } . { file_format } " )
# Send the transcript
if transcript [ ' text ' ] == " " :
transcript [ ' text ' ] = " [Silence] "
2023-04-12 10:42:28 +00:00
2023-04-19 10:53:18 +00:00
if audioMessage and user_data [ " options " ] [ " whisper_to_chat " ] :
chatGPT_response = await messageGPT ( transcript [ ' text ' ] , str ( chat_id ) , update . effective_user . name )
2023-04-12 15:03:09 +00:00
transcript [ ' text ' ] = " > " + transcript [ ' text ' ] + " \n \n " + chatGPT_response
2023-03-22 11:44:08 +00:00
# Check if the transcript length is longer than 4095 characters
if len ( transcript [ ' text ' ] ) > 4095 :
# Split the transcript into multiple messages without breaking words in half
max_length = 4096
words = transcript [ ' text ' ] . split ( )
current_message = " "
for word in words :
if len ( current_message ) + len ( word ) + 1 > max_length :
await context . bot . send_message ( chat_id = chat_id , text = current_message )
current_message = " "
current_message + = f " { word } "
if current_message :
await context . bot . send_message ( chat_id = chat_id , text = current_message )
2023-03-07 20:06:14 +00:00
else :
2023-03-22 11:44:08 +00:00
await context . bot . send_message ( chat_id = chat_id , text = transcript [ ' text ' ] )
2023-04-19 10:53:18 +00:00
# Update user data in the database
database . update_user ( str ( chat_id ) , user_data )
2023-03-06 15:20:37 +00:00
@restricted
async def chat ( update : Update , context : ContextTypes . DEFAULT_TYPE ) :
2023-03-22 11:44:08 +00:00
chat_id = str ( update . effective_chat . id )
2023-03-23 08:14:32 +00:00
await context . bot . send_chat_action ( chat_id = chat_id , action = ChatAction . TYPING )
2023-03-22 11:44:08 +00:00
# Check if replying and add context
2023-03-07 20:06:14 +00:00
if hasattr ( update . message . reply_to_message , " text " ) :
2023-04-12 15:03:09 +00:00
user_prompt = f " In reply to: ' { update . message . reply_to_message . text } ' \n --- \n { update . message . text } "
2023-03-07 20:06:14 +00:00
else :
2023-03-22 11:44:08 +00:00
user_prompt = update . message . text
2023-03-07 20:06:14 +00:00
2023-04-12 10:42:28 +00:00
# Use messageGPT function to get the response
2023-04-19 10:53:18 +00:00
assistant_message = await messageGPT ( user_prompt , chat_id , update . effective_user . name )
2023-04-12 10:42:28 +00:00
await context . bot . send_message ( chat_id = update . effective_chat . id , text = assistant_message )
2023-03-06 15:20:37 +00:00
@restricted
async def clear ( update : Update , context : ContextTypes . DEFAULT_TYPE ) - > None :
2023-04-19 10:53:18 +00:00
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. ' )
2023-03-06 15:20:37 +00:00
@restricted
async def usage ( update : Update , context : ContextTypes . DEFAULT_TYPE ) - > None :
2023-04-19 10:53:18 +00:00
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 : .2 f } % 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 } """
2023-03-06 15:20:37 +00:00
await context . bot . send_message ( chat_id = update . effective_chat . id , text = info_message )
@restricted
async def _help ( update : Update , context : ContextTypes . DEFAULT_TYPE ) - > None :
2023-04-12 10:42:28 +00:00
help_message = """ Here ' s what you can do: \n \n
- / imagine < prompt > to generate an image with DALL - E \n - Send a message to chat with ChatGPT \n
- Send an audio to transcribe to text with Whisper . \n \n
- / settings To change your settings . \n
- / usage To get your usage statistics . \n
- / clear To clear you chatgpt message context ( start a new chat ) . """
2023-03-06 15:20:37 +00:00
await context . bot . send_message ( chat_id = update . effective_chat . id , text = help_message )
2023-04-12 10:42:28 +00:00
# Function to generate the settings buttons
def generate_settings_markup ( chat_id : str ) - > InlineKeyboardMarkup :
keyboard = [
[
InlineKeyboardButton ( " Increase Temperature " , callback_data = f " setting_increase_temperature_ { chat_id } " ) ,
InlineKeyboardButton ( " Decrease Temperature " , callback_data = f " setting_decrease_temperature_ { chat_id } " )
] ,
[
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 } " )
] ,
2023-04-19 10:53:18 +00:00
[
InlineKeyboardButton ( " Enable assistant voice " , callback_data = f " setting_enable_voice_ { chat_id } " ) ,
InlineKeyboardButton ( " Disable assistant voice " , callback_data = f " setting_disable_voice_ { chat_id } " )
] ,
2023-04-12 10:42:28 +00:00
[
InlineKeyboardButton ( " Increase Context " , callback_data = f " setting_increase_context_ { chat_id } " ) ,
InlineKeyboardButton ( " Decrease Context " , callback_data = f " setting_decrease_context_ { chat_id } " )
]
]
return InlineKeyboardMarkup ( keyboard )
@restricted
async def settings ( update : Update , context : ContextTypes . DEFAULT_TYPE ) :
chat_id = update . effective_chat . id
settings_markup = generate_settings_markup ( chat_id )
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 ) :
2023-04-19 10:53:18 +00:00
user_data = await getUserData ( update . effective_chat . id )
2023-04-12 10:42:28 +00:00
query = update . callback_query
action , chat_id = query . data . rsplit ( " _ " , 1 )
2023-04-19 10:53:18 +00:00
# Temperature
2023-04-12 10:42:28 +00:00
if action . startswith ( " setting_increase_temperature " ) :
2023-04-19 10:53:18 +00:00
user_data [ " options " ] [ " temperature " ] = min ( user_data [ " options " ] [ " temperature " ] + 0.1 , 1 )
2023-04-12 10:42:28 +00:00
elif action . startswith ( " setting_decrease_temperature " ) :
2023-04-19 10:53:18 +00:00
user_data [ " options " ] [ " temperature " ] = max ( user_data [ " options " ] [ " temperature " ] - 0.1 , 0 )
# Whisper to GPT
2023-04-12 10:42:28 +00:00
elif action . startswith ( " setting_enable_whisper " ) :
print ( f " enabling whisper for { chat_id } " )
2023-04-19 10:53:18 +00:00
user_data [ " options " ] [ " whisper_to_chat " ] = True
2023-04-12 10:42:28 +00:00
elif action . startswith ( " setting_disable_whisper " ) :
print ( f " disabling whisper for { chat_id } " )
2023-04-19 10:53:18 +00:00
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
2023-04-12 10:42:28 +00:00
elif action . startswith ( " setting_increase_context " ) :
2023-04-19 10:53:18 +00:00
user_data [ " options " ] [ " max-context " ] = min ( user_data [ " options " ] [ " max-context " ] + 1 , MAX_USER_CONTEXT )
2023-04-12 10:42:28 +00:00
elif action . startswith ( " setting_decrease_context " ) :
2023-04-19 10:53:18 +00:00
user_data [ " options " ] [ " max-context " ] = max ( user_data [ " options " ] [ " max-context " ] - 1 , 1 )
2023-04-12 10:42:28 +00:00
settings_markup = generate_settings_markup ( chat_id )
await query . edit_message_text ( text = " Choose a setting option: " , reply_markup = settings_markup )
# Remove the settings message
await context . bot . delete_message ( chat_id = query . message . chat_id , message_id = query . message . message_id )
# Send a message displaying the updated settings
2023-04-19 10:53:18 +00:00
settings_message = f """ Updated settings: \n \n Temperature: { user_data [ ' options ' ] [ ' temperature ' ] } \n Whisper to Chat: { user_data [ ' options ' ] [ ' whisper_to_chat ' ] } \n Assistant voice: { user_data [ ' options ' ] [ ' assistant_voice_chat ' ] } \n Context Length: { user_data [ " options " ] [ " max-context " ] } """
2023-04-12 10:42:28 +00:00
await context . bot . send_message ( chat_id = chat_id , text = settings_message )
2023-03-06 15:20:37 +00:00
if __name__ == ' __main__ ' :
2023-04-19 10:53:18 +00:00
database . init_database ( )
2023-03-06 15:20:37 +00:00
try :
2023-03-22 11:44:08 +00:00
ALLOWED_USERS = os . environ . get ( " BOT_ALLOWED_USERS " ) . split ( " , " )
2023-04-19 10:53:18 +00:00
except ( Exception ) :
2023-03-06 15:20:37 +00:00
ALLOWED_USERS = ALLOWED_USERS
print ( f " Allowed users: { ALLOWED_USERS } " )
print ( f " System prompt: { SYSTEM_PROMPT } " )
2023-03-22 11:44:08 +00:00
application = ApplicationBuilder ( ) . token ( os . environ . get ( " BOT_TOKEN " ) ) . build ( )
2023-03-06 15:20:37 +00:00
start_handler = CommandHandler ( ' start ' , start )
2023-04-12 10:42:28 +00:00
application . add_handler ( start_handler )
2023-03-06 15:20:37 +00:00
clear_handler = CommandHandler ( ' clear ' , clear )
2023-04-12 10:42:28 +00:00
application . add_handler ( clear_handler )
2023-03-06 15:20:37 +00:00
info_handler = CommandHandler ( ' usage ' , usage )
2023-04-12 10:42:28 +00:00
application . add_handler ( info_handler )
2023-03-06 15:20:37 +00:00
help_handler = CommandHandler ( ' help ' , _help )
2023-04-12 10:42:28 +00:00
application . add_handler ( help_handler )
2023-03-06 15:20:37 +00:00
imagine_handler = CommandHandler ( ' imagine ' , imagine )
2023-04-12 10:42:28 +00:00
application . add_handler ( imagine_handler )
2023-03-06 15:20:37 +00:00
2023-04-12 10:42:28 +00:00
settings_handler = CommandHandler ( ' settings ' , settings )
application . add_handler ( settings_handler )
2023-03-06 15:20:37 +00:00
application . add_handler ( MessageHandler ( filters . TEXT & ~ filters . COMMAND , chat ) )
application . add_handler ( MessageHandler ( filters . ATTACHMENT & ~ filters . COMMAND , attachment ) )
2023-04-12 10:42:28 +00:00
settings_callback_handler = CallbackQueryHandler ( settings_callback )
application . add_handler ( settings_callback_handler )
2023-03-06 15:20:37 +00:00
application . run_polling ( )