Skip to content

Commit

Permalink
v0.7504 - fixed usage charts & redone directory mapping
Browse files Browse the repository at this point in the history
  • Loading branch information
FlyingFathead committed Oct 6, 2024
1 parent f1e5937 commit 7e3e2c0
Show file tree
Hide file tree
Showing 3 changed files with 132 additions and 25 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -236,6 +236,7 @@ If you run into any issues, consult the logs or reach out on the repository's [I
---
# Changelog
- v0.7504 - fixed usage logs and charts directory mapping
- v0.7503 - improved message formatting & error catching
- v0.7502 - added `docker_setup.sh` for easier Docker-based deployment
- v0.7501 - `Dockerfile` and better error catching when receiving `401 Unauthorized`
Expand Down
111 changes: 92 additions & 19 deletions src/bot_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,59 +107,132 @@ async def set_system_message_command(update: Update, context: CallbackContext, b
logging.info(f"User {user_id} attempted to set system message but provided no new message.")
await update.message.reply_text("Please provide the new system message in the command line, i.e.: /setsystemmessage My new system message to the AI on what it is, where it is, etc.")


# /usage (admin command)
async def usage_command(update: Update, context: CallbackContext, bot_instance):
# bot_commands.py
async def usage_command(update: Update, context: CallbackContext):
bot_instance = context.bot_data.get('bot_instance') # Retrieve the bot instance from context

if not bot_instance:
await update.message.reply_text("Internal error: Bot instance not found.")
logging.error("Bot instance not found in context.bot_data")
return

logging.info(f"User {update.message.from_user.id} invoked /usage command")

if bot_instance.bot_owner_id == '0':
await update.message.reply_text("The `/usage` command is disabled.")
logging.info("Usage command is disabled until a bot owner is defined in `config.ini`.")
return

if str(update.message.from_user.id) != bot_instance.bot_owner_id:
await update.message.reply_text("You don't have permission to use this command.")
logging.info(f"User {update.message.from_user.id} does not have permission to use /usage")
return

# Define current_date before entering the try block
# Correct path to token_usage.json inside logs/ directory
# token_usage_file = os.path.join(bot_instance.data_directory, 'logs', 'token_usage.json')
token_usage_file = os.path.join(bot_instance.logs_directory, 'token_usage.json')

logging.info(f"Looking for token usage file at: {token_usage_file}")
current_date = datetime.datetime.utcnow()

try:
if os.path.exists(bot_instance.token_usage_file):
with open(bot_instance.token_usage_file, 'r') as file:
if os.path.exists(token_usage_file):
with open(token_usage_file, 'r') as file:
token_usage_history = json.load(file)

# Prune token usage history based on the previously defined current_date
logging.info("Loaded token usage history successfully")

# Prune token usage history
cutoff_date = current_date - datetime.timedelta(days=bot_instance.max_history_days)
token_usage_history = {date: usage for date, usage in token_usage_history.items() if datetime.datetime.strptime(date, '%Y-%m-%d') >= cutoff_date}
token_usage_history = {
date: usage for date, usage in token_usage_history.items()
if datetime.datetime.strptime(date, '%Y-%m-%d') >= cutoff_date
}
logging.info("Pruned token usage history based on cutoff date")
else:
token_usage_history = {}
logging.warning(f"Token usage file does not exist at: {token_usage_file}")
except json.JSONDecodeError:
await update.message.reply_text("Error reading token usage history.")
logging.error("JSONDecodeError while reading token_usage.json")
return
except Exception as e:
await update.message.reply_text(f"An unexpected error occurred: {e}")
logging.error(f"Unexpected error in usage_command: {e}")
return

# Since current_date is now defined outside the try block, it will always be available here
today_usage = token_usage_history.get(current_date.strftime('%Y-%m-%d'), 0)
token_cap_info = f"Today's usage: {today_usage} tokens\n" \
f"Daily token cap: {'No cap' if bot_instance.max_tokens_config == 0 else f'{bot_instance.max_tokens_config} tokens'}\n\n" \
"Token Usage History:\n"
token_cap_info = (
f"Today's usage: {today_usage} tokens\n"
f"Daily token cap: {'No cap' if bot_instance.max_tokens_config == 0 else f'{bot_instance.max_tokens_config} tokens'}\n\n"
"Token Usage History:\n"
)

for date, usage in sorted(token_usage_history.items()):
token_cap_info += f"{date}: {usage} tokens\n"

await update.message.reply_text(token_cap_info)
logging.info("Sent usage information to user")

# /usagechart (admin command)
async def usage_chart_command(update: Update, context: CallbackContext):
bot_instance = context.bot_data.get('bot_instance') # Retrieve the bot instance from context

# /usagechart (admin command, to get chart type usage statistics)
async def usage_chart_command(update: Update, context: CallbackContext, bot_instance, token_usage_file):
if not bot_instance:
await update.message.reply_text("Internal error: Bot instance not found.")
logging.error("Bot instance not found in context.bot_data")
return

logging.info(f"User {update.message.from_user.id} invoked /usagechart command")

if bot_instance.bot_owner_id == '0':
await update.message.reply_text("The `/usagechart` command is disabled.")
logging.info("Usagechart command is disabled")
return

if str(update.message.from_user.id) != bot_instance.bot_owner_id:
await update.message.reply_text("You don't have permission to use this command.")
logging.info(f"User {update.message.from_user.id} does not have permission to use /usagechart")
return

output_image_file = 'token_usage_chart.png'
generate_usage_chart(token_usage_file, output_image_file)

with open(output_image_file, 'rb') as file:
await context.bot.send_photo(chat_id=update.message.chat_id, photo=file)

# Define paths
token_usage_file = os.path.join(bot_instance.logs_directory, 'token_usage.json')
output_image_file = os.path.join(bot_instance.data_directory, 'token_usage_chart.png')

logging.info(f"Looking for token usage file at: {token_usage_file}")
logging.info(f"Output image file will be at: {output_image_file}")

# Ensure the data directory exists
try:
if not os.path.exists(bot_instance.data_directory):
os.makedirs(bot_instance.data_directory, exist_ok=True)
bot_instance.logger.info(f"Created data directory at {bot_instance.data_directory}")
except OSError as e:
bot_instance.logger.error(f"Failed to create data directory {bot_instance.data_directory}: {e}")
await update.message.reply_text(f"Failed to create the data directory for the chart. Please check the bot's permissions.")
return

# Generate the usage chart
try:
generate_usage_chart(token_usage_file, output_image_file)
bot_instance.logger.info(f"Generated usage chart at {output_image_file}")
except Exception as e:
bot_instance.logger.error(f"Failed to generate usage chart: {e}")
await update.message.reply_text("Failed to generate usage chart.")
return

# Try to open and send the generated chart image
try:
with open(output_image_file, 'rb') as file:
await context.bot.send_photo(chat_id=update.message.chat_id, photo=file)
bot_instance.logger.info(f"Sent usage chart to chat_id {update.message.chat_id}")
except FileNotFoundError:
await update.message.reply_text("Token usage chart not found. Please ensure it's being generated correctly.")
bot_instance.logger.error("Token usage chart file not found: %s", output_image_file)
except Exception as e:
await update.message.reply_text("Failed to send the usage chart.")
bot_instance.logger.error(f"Error sending usage chart: {e}")

# /reset
async def reset_command(update: Update, context: CallbackContext, bot_owner_id, reset_enabled, admin_only_reset):
Expand Down
45 changes: 39 additions & 6 deletions src/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
# https://github.com/FlyingFathead/TelegramBot-OpenAI-API
#
# version of this program
version_number = "0.7503"
version_number = "0.7504"

# Add the project root directory to Python's path
import sys
Expand Down Expand Up @@ -147,9 +147,35 @@ def load_config(self):
self.enable_whisper = self.config.getboolean('EnableWhisper', True)
self.max_voice_message_length = self.config.getint('MaxDurationMinutes', 5)

self.data_directory = self.config.get('DataDirectory', 'data') # Default to 'data' if not set
self.data_directory = self.config.get('DataDirectory', 'data') # Default to 'data' if not set
self.max_storage_mb = self.config.getint('MaxStorageMB', 100) # Default to 100 MB if not set

# set directories
# Get the project root
project_root = Path(__file__).resolve().parents[1]

# Combine project root with the relative data directory from config
self.data_directory = str(project_root / self.config.get('DataDirectory', 'data'))

# Ensure the data directory exists or handle creation failure
try:
if not os.path.exists(self.data_directory):
os.makedirs(self.data_directory, exist_ok=True)
self.logger.info(f"Created data directory at {self.data_directory}")
except OSError as e:
self.logger.error(f"Failed to create data directory {self.data_directory}: {e} -- Some commands might be disabled due to this.")

# set the logs directory
self.logs_directory = str(project_root / self.config.get('LogsDirectory', 'logs'))

# Ensure the logs directory exists or handle creation failure
try:
if not os.path.exists(self.logs_directory):
os.makedirs(self.logs_directory, exist_ok=True)
self.logger.info(f"Created logs directory at {self.logs_directory}")
except OSError as e:
self.logger.error(f"Failed to create logs directory {self.logs_directory}: {e} -- Some commands might be disabled due to this.")

self.logfile_enabled = self.config.getboolean('LogFileEnabled', True)
self.logfile_file = self.config.get('LogFile', 'bot.log')
self.chat_logging_enabled = self.config.getboolean('ChatLoggingEnabled', False)
Expand Down Expand Up @@ -304,7 +330,11 @@ def error(self, update: Update, context: CallbackContext) -> None:
def run(self):
application = Application.builder().token(self.telegram_bot_token).build()
application.get_updates_read_timeout = self.timeout


# Store bot_instance in bot_data for access in handlers
application.bot_data['bot_instance'] = self
self.logger.info("Stored bot_instance in context.bot_data")

# Text handler
application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message))

Expand All @@ -325,9 +355,12 @@ def run(self):
# application.add_handler(CommandHandler("restart", partial(bot_commands.restart_command, bot_owner_id=self.bot_owner_id)))
# application.add_handler(CommandHandler("usage", partial(bot_commands.usage_command, bot_owner_id=self.bot_owner_id, total_token_usage=self.total_token_usage, max_tokens_config=self.max_tokens_config)))
# application.add_handler(CommandHandler("updateconfig", partial(bot_commands.update_config_command, bot_owner_id=self.bot_owner_id)))
application.add_handler(CommandHandler("usagechart", partial(bot_commands.usage_chart_command, bot_instance=self, token_usage_file='token_usage.json')))
application.add_handler(CommandHandler("usage", partial(bot_commands.usage_command, bot_instance=self)))

# application.add_handler(CommandHandler("usagechart", partial(bot_commands.usage_chart_command, bot_instance=self, token_usage_file='token_usage.json')))
# application.add_handler(CommandHandler("usage", partial(bot_commands.usage_command, bot_instance=self)))

application.add_handler(CommandHandler("usagechart", bot_commands.usage_chart_command))
application.add_handler(CommandHandler("usage", bot_commands.usage_command))

# Register the /reset command
application.add_handler(CommandHandler("reset", partial(bot_commands.reset_command,
bot_owner_id=self.bot_owner_id,
Expand Down

0 comments on commit 7e3e2c0

Please sign in to comment.