Added better logging

.gitignore

@@ -6,3 +6,5 @@ __pycache__/
 !/conf/example-workflow_settings.json
 /log.txt
 /log.txt.*
+/*.log.txt
+/*.log.txt.*

bot.py

@@ -4,9 +4,8 @@ import lib.log as log
 from lib.settings import *
 from lib.events import *
 
-intents     = discord.Intents(messages=True, guilds=True, message_content=True, reactions=True)
+intents = discord.Intents(messages=True, guilds=True, message_content=True, reactions=True)
-client      = discord.Client(intents=intents)
+client  = discord.Client(intents=intents)
-settings    = {}
 
 @client.event
 async def on_message(msg):
@@ -14,23 +13,16 @@ async def on_message(msg):
 
 @client.event
 async def on_raw_reaction_add(rxn):
-    ALLOW_REPEAT    = settings["allow_repeat"]
+    await on_message_or_reaction(client, rxn)
-    REPEAT_EMOJI    = settings["repeat_emoji"]
-
-    if ALLOW_REPEAT:
-        if rxn.emoji.name == REPEAT_EMOJI:
-            await on_message_or_reaction(client, rxn)
 
 @client.event
 async def on_ready():
-    log.write("READY.")
+    log.write("READY.", log.colors.fg.lightgreen)
 
 def main():
-    global settings
+    log.write("Starting up...", log.colors.fg.yellow)
     settings = get_settings(initialize=True)
     client.run(settings["token"], log_handler=None)
 
 if __name__ == "__main__":
-    log.write("test")
     main()
-

lib/events.py

@@ -13,6 +13,8 @@ from lib.parser import *
 from lib.comfyui import *
 
 async def on_message_or_reaction(client, obj):
+    global RUN_TIME
+
     msg     = None
     chl     = None
     user    = None
@@ -56,7 +58,6 @@ async def on_message_or_reaction(client, obj):
         chl_topic_parts     = chl.topic.split(",")
         chl_topic_part_1    = chl_topic_parts[0]
 
-
     # Try different paths to find workflow .json file:
     workflow_paths = [
         f"conf/{chl.category.name}/{chl_topic_part_1}",
@@ -69,17 +70,11 @@ async def on_message_or_reaction(client, obj):
     using_workflow_path = None
     using_settings_path = None
 
-    log.write("\nSearching for workflow in order:\n")
-
     for path in workflow_paths:
-        check_text  = " "
+        try_path = f"{path}.json"
-        try_path    = f"{path}.json"
-
         if os.path.isfile(try_path):
             using_workflow_path = try_path
-            check_text          = "x"
+            log.write(f"Workflow {log.colors.fg.lightcyan}{try_path}")
-
-        log.write(f"[ {check_text} ] Selecting: {try_path}")
 
     if using_workflow_path is None:
         return
@@ -103,16 +98,11 @@ async def on_message_or_reaction(client, obj):
 
     settings = get_settings()
 
-    log.write("Merging settings in order:")
     for path in setting_paths:
-        check_text  = " "
+        try_path = f"{path}_settings.json"
-        try_path    = f"{path}_settings.json"
-        
         if os.path.isfile(try_path):
-            settings    = merge_dicts(settings, read_json(try_path, {}))
+            settings = merge_dicts(settings, read_json(try_path, {}))
-            check_text  = "x"
+            log.write(f"Merging {log.colors.fg.lightcyan}{try_path}")
-
-        log.write(f"[ {check_text} ] Merging: {try_path}")
 
     #
     #
@@ -132,6 +122,7 @@ async def on_message_or_reaction(client, obj):
         allowed_users = settings["allowed_users"]
         if isinstance(allowed_users, list):
             if str(user.id) not in allowed_users and str(user.name) not in allowed_users:
+                log.write(f"User {user.name} or their ID not in list of allowed users", log.colors.fg.lightred)
                 return
 
     # If ignored_users is specified, ignore listed users:
@@ -139,9 +130,9 @@ async def on_message_or_reaction(client, obj):
         ignored_users = settings["ignored_users"]
         if isinstance(ignored_users, list):
             if str(user.id) in ignored_users or str(user.name) in ignored_users:
+                log.write(f"User {user.name} or their ID is explicitly being ignored", log.colors.fg.lightred)
                 return
 
-
     # Read the found .json file:
     workflow_json = ""
     with open(using_workflow_path, "r") as file:
@@ -150,93 +141,103 @@ async def on_message_or_reaction(client, obj):
     # Break the user message into prompt parameters:
     params = get_prompt_parameters(msg.content, settings)
 
-
-    max_key_size = 0
     for k in params.keys():
-        new_key_size = len(k) + 4
+        v       = params[k]
-        if new_key_size > max_key_size:
+        k       = k.upper()
-            max_key_size = new_key_size
-
-    log.write("Replacing in workflow JSON body:")
-    for k in params.keys():
-        v = params[k]
-        k = k.upper()
-
         label   = f"__{k}__"
-        while len(label) < max_key_size:
+        log.write(f"Replacing {log.colors.fg.yellow}{label}{log.colors.reset} with {log.colors.fg.white}{v}")
-            label = label + " "
 
-        log.write(f"[ - ] Replacing: {label} with {v}")
+    log.write_prompt(client, msg, chl, user, author, roles, rxn)
 
+    try:
+        time_start = time.perf_counter()
 
-    time_start = time.perf_counter()
+        # Indicate to the user something is happening:
+        await msg.add_reaction(WAITING_EMOJI)
+        await chl.typing()
 
-    # Indicate to the user something is happening:
+        repeat_n_times  = 1 if "repeat_n_times" not in params.keys() else params["repeat_n_times"]
-    await msg.add_reaction(WAITING_EMOJI)
+        all_attachments = []
-    await chl.typing()
 
-    repeat_n_times  = 1 if "repeat_n_times" not in params.keys() else params["repeat_n_times"]
+        for i in range(0, repeat_n_times):
-    all_attachments = []
+            workflow_json_clone = workflow_json
+            params_clone        = params.copy()
 
-    for i in range(0, repeat_n_times):
+            # If the seed is not specified, generate one:
-        workflow_json_clone = workflow_json
+            if "seed" not in params.keys() or params["seed"] is None:
-        params_clone        = params.copy()
+                new_seed                = random.randint(1, 999_999_999_999_999)
+                params_clone["seed"]    = new_seed
+                log.write(f"Replacing {log.colors.fg.yellow}__SEED__{log.colors.reset} with {log.colors.fg.white}{new_seed}")
 
-        # If the seed is not specified, generate one:
+            # Make replacements, e.g.
-        if "seed" not in params.keys() or params["seed"] is None:
+            # __WIDTH__ to value of params["width"]
-            new_seed                = random.randint(1, 999_999_999_999_999)
+            # __POSITIVE__ to value of params["positive"]
-            params_clone["seed"]    = new_seed
+            for k in params_clone.keys():
+                v = params_clone[k]
+                k = k.upper()
 
-            label   = f"__SEED__"
+                if v is None:
-            while len(label) < max_key_size:
+                    continue
-                label = label + " "
 
-            log.write(f"[ - ] Replacing: {label} with {new_seed}")
+                workflow_json_clone = re.sub(rf"__{k}__", str(v), workflow_json_clone)
 
-        # Make replacements, e.g.
+            # Must be valid JSON:
-        # __WIDTH__ to value of params["width"]
+            workflow        = json.loads(workflow_json_clone)
-        # __POSITIVE__ to value of params["positive"]
+            all_attachments = all_attachments + await get_comfyui_generations(settings["api_url"], workflow)
-        for k in params_clone.keys():
-            v = params_clone[k]
-            k = k.upper()
 
-            if v is None:
+        bytes_generated     = 0
-                continue
+        images_generated    = 0
 
-            workflow_json_clone = re.sub(rf"__{k}__", str(v), workflow_json_clone)
+        # Process 8 attachments at a time per one Discord message:
+        while True:
+            attachments_buffer  = all_attachments[:8]
+            discord_files       = []
+            
+            for a in attachments_buffer:
+                bytes_generated     = bytes_generated + len(a)
+                images_generated    = images_generated + 1
+                discord_file        = discord.File(io.BytesIO(a), filename="file.png")
+                discord_files.append(discord_file)
 
-        # Must be valid JSON:
+            if len(discord_files) < 1: 
-        workflow        = json.loads(workflow_json_clone)
+                discord_files = None
-        all_attachments = all_attachments + await get_comfyui_generations(settings["api_url"], workflow)
 
-    # Process 8 attachments at a time per one Discord message:
+            post = await chl.send(files=discord_files, content=msg.content, reference=msg)
-    while True:
+            del all_attachments[:8]
-        attachments_buffer  = all_attachments[:8]
+
-        discord_files       = []
+            if len(all_attachments) < 1:
+                break
+
+        if ALLOW_REPEAT:
+            if SHOW_REPEAT:
+                await msg.add_reaction(REPEAT_EMOJI)
+                await post.add_reaction(REPEAT_EMOJI)
         
-        for a in attachments_buffer:
+        await msg.remove_reaction(WAITING_EMOJI, client.user)
-            #base64_str          = a
-            #base64_bytes        = base64_str.encode("ascii")
-            #attachment_bytes    = base64.b64decode(base64_bytes)
 
-            f = discord.File(io.BytesIO(a), filename="file.png")
+        time_end    = time.perf_counter()
-            discord_files.append(f)
+        time_taken  = time_end - time_start
 
-        if len(discord_files) < 1: 
+        log.write(f"{time_taken:0.2f}s taken.", log.colors.fg.lightgreen)
-            discord_files = None
 
-        post = await chl.send(files=discord_files, content=msg.content, reference=msg)
+        #
-        del all_attachments[:8]
+        #
+        #
 
-        if len(all_attachments) < 1:
+        log.write_stat(f"user.{user.name}.time_taken", f"{time_taken:0.2f}")
-            break
+        log.write_stat(f"user.{user.name}.bytes_generated", f"{bytes_generated}")
+        log.write_stat(f"user.{user.name}.images_generated", f"{images_generated}")
 
-    if ALLOW_REPEAT:
-        if SHOW_REPEAT:
-            await msg.add_reaction(REPEAT_EMOJI)
-            await post.add_reaction(REPEAT_EMOJI)
-    
-    await msg.remove_reaction(WAITING_EMOJI, client.user)
 
-    time_end    = time.perf_counter()
+        t_time_taken = log.get_stat(r".*\.time_taken\: (.+)$")
-    time_taken  = time_end - time_start
+        t_bytes_generated = log.get_stat(r".*\.bytes_generated\: (.+)$")
+        t_images_generated = log.get_stat(r".*\.images_generated\: (.+)$")
+
+        kb_gen = t_bytes_generated / 1024
+        mb_gen = kb_gen / 1024
+
+        await client.change_presence(activity=discord.Game(name=f"{t_time_taken:0.1f}s, {mb_gen:0.1f} MiB, {t_images_generated:0.0f}x"))
+
+
+    finally:
+        lock = False

lib/log.py

@@ -1,41 +1,121 @@
 import sys
 import logging
 import logging.handlers
+import re
 
 LOG_SETUP = False
 
+# The colors class below is from:
+# https://www.geeksforgeeks.org/print-colors-python-terminal
+class colors:
+    reset           = "\033[0m"
+    bold            = "\033[01m"
+    disable         = "\033[02m"
+    underline       = "\033[04m"
+    reverse         = "\033[07m"
+    strikethrough   = "\033[09m"
+    invisible       = "\033[08m"
+     
+    class fg:
+        white       = "\033[97m"
+        black       = "\033[30m"
+        red         = "\033[31m"
+        green       = "\033[32m"
+        orange      = "\033[33m"
+        blue        = "\033[34m"
+        purple      = "\033[35m"
+        cyan        = "\033[36m"
+        lightgrey   = "\033[37m"
+        darkgrey    = "\033[90m"
+        lightred    = "\033[91m"
+        lightgreen  = "\033[92m"
+        yellow      = "\033[93m"
+        lightblue   = "\033[94m"
+        pink        = "\033[95m"
+        lightcyan   = "\033[96m"
+
+    class bg:
+        black       = "\033[40m"
+        red         = "\033[41m"
+        green       = "\033[42m"
+        orange      = "\033[43m"
+        blue        = "\033[44m"
+        purple      = "\033[45m"
+        cyan        = "\033[46m"
+        lightgrey   = "\033[47m"
+
+class PlainTextFormatter(logging.Formatter):
+    def format(self, record):
+        log_message = super().format(record)
+        log_message = re.sub(r"\033\[[0-9]{1,2}m", "", log_message)
+        return log_message
+
 # Make it easy to log text:
-def write(input_str):
+def write(input_str, color=colors.fg.lightgrey, logger_name="discord.bot"):
     global LOG_SETUP
+
     if not LOG_SETUP:
+        logging.getLogger("discord").setLevel(logging.DEBUG)
+        logging.getLogger("discord.http").setLevel(logging.INFO)
 
-        logger = logging.getLogger('discord')
+        line_format         = "[{asctime}] [{threadName}] [{levelname}] {name}: {message}"
-        logger.setLevel(logging.DEBUG)
+        timestamp_format    = "%Y-%m-%d %H:%M:%S"
-        logging.getLogger('discord.http').setLevel(logging.INFO)
+        console_formatter   = logging.Formatter(colors.fg.darkgrey + line_format, timestamp_format, style="{")
+        file_formatter      = PlainTextFormatter(line_format, timestamp_format, style="{")
+        prompt_formatter    = PlainTextFormatter("[{asctime}] {message}", timestamp_format, style="{")
+        stats_formatter     = PlainTextFormatter("{message}", timestamp_format, style="{")
+        
+        file_handler = logging.handlers.RotatingFileHandler(filename="log.txt", encoding="utf-8", maxBytes=32 * 1024 * 1024, backupCount=10)
+        file_handler.setFormatter(file_formatter)
 
+        console_handler = logging.StreamHandler()
+        console_handler.setFormatter(console_formatter)
 
-        logFormatter = logging.Formatter(
+        root_logger = logging.getLogger()
-            "[{asctime}] [{threadName:<16}] [{levelname:<8}] {name}: {message}",
+        root_logger.addHandler(file_handler)
-            "%Y-%m-%d %H:%M:%S",
+        root_logger.addHandler(console_handler)
-            style="{")
 
-        rootLogger = logging.getLogger()
+        prompt_logger   = logging.getLogger("discord.bot.prompts")
+        prompt_handler  = logging.handlers.RotatingFileHandler(filename="prompts.log.txt", encoding="utf-8", maxBytes=32 * 1024 * 1024, backupCount=10)
+        prompt_handler.setFormatter(prompt_formatter)
+        prompt_logger.addHandler(prompt_handler)
 
-        #fileHandler = logging.FileHandler("log.txt")
+        stats_logger   = logging.getLogger("discord.bot.stats")
-        fileHandler = logging.handlers.RotatingFileHandler(
+        stats_handler  = logging.handlers.RotatingFileHandler(filename="stats.log.txt", encoding="utf-8", maxBytes=32 * 1024 * 1024, backupCount=10)
-            filename="log.txt",
+        stats_handler.setFormatter(stats_formatter)
-            encoding="utf-8",
+        stats_logger.addHandler(stats_handler)
-            maxBytes=32 * 1024 * 1024,
-            backupCount=10,
-        )
-
-        fileHandler.setFormatter(logFormatter)
-        rootLogger.addHandler(fileHandler)
-
-        consoleHandler = logging.StreamHandler()
-        consoleHandler.setFormatter(logFormatter)
-        rootLogger.addHandler(consoleHandler)
 
         LOG_SETUP = True
 
-    logging.getLogger("discord").info(input_str)
+    logging.getLogger(logger_name).info(f"{color}{input_str}{colors.reset}")
+
+def write_prompt(client=None, msg=None, chl=None, user=None, author=None, roles=None, rxn=None):
+    bot_name        = client.user.name + "#" + client.user.discriminator
+    server_name     = msg.channel.guild.name
+    category_name   = msg.channel.category.name
+    channel_name    = msg.channel.name
+    user_name       = user.name
+
+    discrim = 0
+    if user.discriminator is not None:
+        discrim = user.discriminator
+
+    user_name       = f"{user_name}#{discrim}"
+    message_content = msg.content
+    message_text    = f"[{bot_name}/{server_name}/{category_name}/{channel_name}/{user_name}] {colors.fg.white}{message_content}"
+
+    write(message_text, logger_name="discord.bot.prompts")
+
+def write_stat(key, value):
+    write(f"{key}: {value}", logger_name="discord.bot.stats")
+
+def get_stat(key, sum=True):
+    value = 0
+    with open("stats.log.txt") as file:
+        lines = [line.strip() for line in file]
+        for line in lines:
+            match = re.search(key, line)
+            if match is not None:
+                value = value + float(match.group(1).strip())
+
+    return value

lib/parser.py

@@ -3,8 +3,9 @@ import re
 import copy
 
 def get_prompt_parameters(user_message, settings):
-    params          = {}
+    params = {}
-    params          = settings["defaults"].copy()
+    if "defaults" in settings.keys():
+        params = settings["defaults"].copy()
 
     mutable_string = user_message
     mutable_string = re.sub(r"[^A-Za-z0-9 \-,\.>\(\)]", "", mutable_string)