Improve logging for normal operation

app.py

@@ -5,8 +5,8 @@ from openai import OpenAI
 import gradio as gr
 
 from theme import apriel
-from utils import COMMUNITY_POSTFIX_URL, get_model_config, log_message, check_format, models_config, \
-    logged_event_handler, DEBUG_MODEL
+from utils import COMMUNITY_POSTFIX_URL, get_model_config, check_format, models_config, \
+    logged_event_handler, DEBUG_MODEL, log_debug, log_info, log_error
 from log_chat import log_chat
 
 MODEL_TEMPERATURE = 0.8
@@ -15,8 +15,6 @@ DEFAULT_OPT_OUT_VALUE = False
 
 DEFAULT_MODEL_NAME = "Apriel-Nemotron-15b-Thinker" if not DEBUG_MODEL else "Apriel-5b"
 
-print(f"Gradio version: {gr.__version__}")
-
 BUTTON_ENABLED = gr.update(interactive=True)
 BUTTON_DISABLED = gr.update(interactive=False)
 INPUT_ENABLED = gr.update(interactive=True)
@@ -37,7 +35,7 @@ openai_client = None
 def app_loaded(state, request: gr.Request):
     message_html = setup_model(DEFAULT_MODEL_NAME, intial=False)
     state['session'] = request.session_hash if request else uuid4().hex
-    log_message(f"app_loaded() --> Session: {state['session']}")
+    log_debug(f"app_loaded() --> Session: {state['session']}")
     return state, message_html
 
 
@@ -50,7 +48,7 @@ def update_model_and_clear_chat(model_name):
 def setup_model(model_name, intial=False):
     global model_config, openai_client
     model_config = get_model_config(model_name)
-    log_message(f"update_model() --> Model config: {model_config}")
+    log_debug(f"update_model() --> Model config: {model_config}")
     openai_client = OpenAI(
         api_key=model_config.get('AUTH_TOKEN'),
         base_url=model_config.get('VLLM_API_URL')
@@ -60,7 +58,7 @@ def setup_model(model_name, intial=False):
     _link = f"<a href='{model_config.get('MODEL_HF_URL')}{COMMUNITY_POSTFIX_URL}' target='_blank'>{_model_hf_name}</a>"
     _description = f"We'd love to hear your thoughts on the model. Click here to provide feedback - {_link}"
 
-    print(f"Switched to model {_model_hf_name}")
+    log_debug(f"Switched to model {_model_hf_name}")
 
     if intial:
         return
@@ -95,14 +93,20 @@ def run_chat_inference(history, message, state):
     state["is_streaming"] = True
     state["stop_flag"] = False
     error = None
+    model_name = model_config.get('MODEL_NAME')
 
     if len(history) == 0:
         state["chat_id"] = uuid4().hex
 
+    if openai_client is None:
+        log_info("Client UI is stale, letting user know to refresh the page")
+        gr.Warning("Client UI is stale, please refresh the page")
+        return history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state
+
     # outputs: model_dropdown, user_input, send_btn, stop_btn, clear_btn, session_state
-    log_message(f"{'-' * 80}")
-    log_message(f"chat_fn() --> Message: {message}")
-    log_message(f"chat_fn() --> History: {history}")
+    log_debug(f"{'-' * 80}")
+    log_debug(f"chat_fn() --> Message: {message}")
+    log_debug(f"chat_fn() --> History: {history}")
 
     try:
         # Check if the message is empty
@@ -112,16 +116,16 @@ def run_chat_inference(history, message, state):
             return history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state
 
         chat_start_count = chat_start_count + 1
-        print(
-            f"{datetime.datetime.now()}: chat_start_count: {chat_start_count}, turns: {int(len(history if history else []) / 3)}")
+        log_info(f"chat_start_count: {chat_start_count}, turns: {int(len(history if history else []) / 3)}, "
+                 f"model: {model_name}")
 
         is_reasoning = model_config.get("REASONING")
 
         # Remove any assistant messages with metadata from history for multiple turns
-        log_message(f"Initial History: {history}")
+        log_debug(f"Initial History: {history}")
         check_format(history, "messages")
         history.append({"role": "user", "content": message})
-        log_message(f"History with user message: {history}")
+        log_debug(f"History with user message: {history}")
         check_format(history, "messages")
 
         # Create the streaming response
@@ -131,18 +135,18 @@ def run_chat_inference(history, message, state):
                                         item.get("role") == "assistant" and
                                         isinstance(item.get("metadata"), dict) and
                                         item.get("metadata", {}).get("title") is not None)]
-            log_message(f"Updated History: {history_no_thoughts}")
+            log_debug(f"Updated History: {history_no_thoughts}")
             check_format(history_no_thoughts, "messages")
-            log_message(f"history_no_thoughts with user message: {history_no_thoughts}")
+            log_debug(f"history_no_thoughts with user message: {history_no_thoughts}")
 
             stream = openai_client.chat.completions.create(
-                model=model_config.get('MODEL_NAME'),
+                model=model_name,
                 messages=history_no_thoughts,
                 temperature=MODEL_TEMPERATURE,
                 stream=True
             )
         except Exception as e:
-            print(f"Error: {e}")
+            log_error(f"Error: {e}")
             error = str(e)
             yield ([{"role": "assistant",
                      "content": "😔 The model is unavailable at the moment. Please try again later."}],
@@ -150,14 +154,14 @@ def run_chat_inference(history, message, state):
             if state["opt_out"] is not True:
                 log_chat(chat_id=state["chat_id"],
                          session_id=state["session"],
-                         model_name=model_config.get('MODEL_NAME'),
+                         model_name=model_name,
                          prompt=message,
                          history=history,
                          info={"is_reasoning": model_config.get("REASONING"), "temperature": MODEL_TEMPERATURE,
                                "stopped": True, "error": str(e)},
                          )
             else:
-                log_message(f"User opted out of chat history. Not logging chat.")
+                log_info(f"User opted out of chat history. Not logging chat. model: {model_name}")
             return history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state
 
         if is_reasoning:
@@ -166,21 +170,21 @@ def run_chat_inference(history, message, state):
                 content="Thinking...",
                 metadata={"title": "🧠 Thought"}
             ))
-            log_message(f"History added thinking: {history}")
+            log_debug(f"History added thinking: {history}")
             check_format(history, "messages")
         else:
             history.append(gr.ChatMessage(
                 role="assistant",
                 content="",
             ))
-            log_message(f"History added empty assistant: {history}")
+            log_debug(f"History added empty assistant: {history}")
             check_format(history, "messages")
 
         output = ""
         completion_started = False
         for chunk in stream:
             if state["stop_flag"]:
-                log_message(f"chat_fn() --> Stopping streaming...")
+                log_debug(f"chat_fn() --> Stopping streaming...")
                 break  # Exit the loop if the stop flag is set
             # Extract the new content from the delta field
             content = getattr(chunk.choices[0].delta, "content", "")
@@ -224,28 +228,31 @@ def run_chat_inference(history, message, state):
             # log_message(f"Yielding messages: {history}")
             yield history, INPUT_DISABLED, SEND_BUTTON_DISABLED, STOP_BUTTON_ENABLED, BUTTON_DISABLED, state
 
-        log_message(f"Final History: {history}")
+        log_debug(f"Final History: {history}")
         check_format(history, "messages")
         yield history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state
     finally:
         if error is None:
-            log_message(f"chat_fn() --> Finished streaming. {chat_start_count} chats started.")
+            log_debug(f"chat_fn() --> Finished streaming. {chat_start_count} chats started.")
             if state["opt_out"] is not True:
                 log_chat(chat_id=state["chat_id"],
                          session_id=state["session"],
-                         model_name=model_config.get('MODEL_NAME'),
+                         model_name=model_name,
                          prompt=message,
                          history=history,
                          info={"is_reasoning": model_config.get("REASONING"), "temperature": MODEL_TEMPERATURE,
                                "stopped": state["stop_flag"]},
                          )
+
             else:
-                log_message(f"User opted out of chat history. Not logging chat.")
+                log_info(f"User opted out of chat history. Not logging chat. model: {model_name}")
         state["is_streaming"] = False
         state["stop_flag"] = False
         return history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state
 
 
+log_info(f"Gradio version: {gr.__version__}")
+
 title = None
 description = None
 theme = apriel
@@ -371,3 +378,4 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
     )
 
 demo.queue(default_concurrency_limit=2).launch(ssr_mode=False, show_api=False)
+log_info("Gradio app launched")

log_chat.py

@@ -10,7 +10,7 @@ from gradio import ChatMessage
 from huggingface_hub import HfApi, hf_hub_download
 
 from timer import Timer
-from utils import log_warning, log_message
+from utils import log_warning, log_info, log_debug, log_error
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 DATASET_REPO_ID = os.environ.get("APRIEL_PROMPT_DATASET")
@@ -18,6 +18,7 @@ CSV_FILENAME = "train.csv"
 
 
 def log_chat(chat_id: str, session_id: str, model_name: str, prompt: str, history: list[str], info: dict) -> None:
+    log_info(f"log_chat() called for chat: {chat_id}, queue size: {log_chat_queue.qsize()}, model: {model_name}")
     log_chat_queue.put((chat_id, session_id, model_name, prompt, history, info))
 
 
@@ -28,12 +29,13 @@ def _log_chat_worker():
             try:
                 _log_chat(chat_id, session_id, model_name, prompt, history, info)
             except Exception as e:
-                log_warning(f"Error logging chat: {e}")
+                log_error(f"Error logging chat: {e}")
         finally:
             log_chat_queue.task_done()
 
 
 def _log_chat(chat_id: str, session_id: str, model_name: str, prompt: str, history: list[str], info: dict) -> bool:
+    log_info(f"_log_chat() storing chat {chat_id}")
     if DATASET_REPO_ID is None:
         log_warning("No dataset repo ID provided. Skipping logging of prompt.")
         return False
@@ -50,9 +52,9 @@ def _log_chat(chat_id: str, session_id: str, model_name: str, prompt: str, histo
     # Check if the dataset repo exists, if not, create it
     try:
         repo_info = api.repo_info(repo_id=DATASET_REPO_ID, repo_type="dataset")
-        log_message(f"log_chat() --> Dataset repo found: {repo_info.id} private={repo_info.private}")
+        log_debug(f"log_chat() --> Dataset repo found: {repo_info.id} private={repo_info.private}")
     except Exception:  # Create new dataset if none exists
-        log_message(f"log_chat() --> No dataset repo found, creating a new one...")
+        log_debug(f"log_chat() --> No dataset repo found, creating a new one...")
         api.create_repo(repo_id=DATASET_REPO_ID, repo_type="dataset", private=True)
 
     # Ensure messages are in the correct format
@@ -97,7 +99,7 @@ def _log_chat(chat_id: str, session_id: str, model_name: str, prompt: str, histo
             )
             pd.read_csv(csv_path)
             file_exists = True
-            log_message(f"log_chat() --> Downloaded existing CSV with {len(pd.read_csv(csv_path))} rows")
+            log_debug(f"log_chat() --> Downloaded existing CSV with {len(pd.read_csv(csv_path))} rows")
             break  # Success, exit the loop
         except Exception as e:
             retry_count += 1
@@ -127,10 +129,10 @@ def _log_chat(chat_id: str, session_id: str, model_name: str, prompt: str, histo
             dump_hub_csv()
             file_exists = False
         else:
-            log_message(f"log_chat() --> CSV {csv_path} has expected headers: {existing_headers}")
+            log_debug(f"log_chat() --> CSV {csv_path} has expected headers: {existing_headers}")
 
     # Write out the new row to the CSV file (append isn't working in HF container, so recreate each time)
-    log_message(f"log_chat() --> Writing CSV file, file_exists={file_exists}")
+    log_debug(f"log_chat() --> Writing CSV file, file_exists={file_exists}")
     try:
         with open(CSV_FILENAME, "w", newline="\n") as f:
             writer = csv.DictWriter(f, fieldnames=new_row.keys())
@@ -140,10 +142,10 @@ def _log_chat(chat_id: str, session_id: str, model_name: str, prompt: str, histo
                     writer.writerow(row.to_dict())  # Write existing rows
             writer.writerow(new_row)  # Write the new row
 
-        log_message("log_chat() --> Wrote out CSV with new row")
+        log_debug("log_chat() --> Wrote out CSV with new row")
         # dump_local_csv()
     except Exception as e:
-        log_warning(f"log_chat() --> Error writing to CSV: {e}")
+        log_error(f"log_chat() --> Error writing to CSV: {e}")
         return False
 
     # Upload updated CSV
@@ -156,8 +158,8 @@ def _log_chat(chat_id: str, session_id: str, model_name: str, prompt: str, histo
     )
     log_timer.add_step("Uploaded updated CSV")
     log_timer.end()
-    log_message("log_chat() --> Finished logging chat")
-    log_message(log_timer.formatted_result())
+    log_debug("log_chat() --> Finished logging chat")
+    log_debug(log_timer.formatted_result())
 
     return True
 
@@ -172,23 +174,23 @@ def dump_hub_csv():
             token=HF_TOKEN  # Only needed if not already logged in
         )
         df = pd.read_csv(csv_path)
-        print(df)
+        log_info(df)
         if (df.empty):
             # show raw contents of downloaded csv file
-            print("Raw file contents:")
+            log_info("Raw file contents:")
             with open(csv_path, 'r') as f:
                 print(f.read())
     except Exception as e:
-        print(f"Error loading CSV from hub: {e}")
+        log_error(f"Error loading CSV from hub: {e}")
 
 
 def dump_local_csv():
     # Verify the file contents by loading it from the local file and printing it out
     try:
         df = pd.read_csv(CSV_FILENAME)
-        print(df)
+        log_info(df)
     except Exception as e:
-        print(f"Error loading CSV from local file: {e}")
+        log_error(f"Error loading CSV from local file: {e}")
 
 
 def test_log_chat():
@@ -207,21 +209,21 @@ def test_log_chat():
                ]
     info = {"additional_info": "Some extra data"}
 
-    log_message("Starting test_log_chat()")
+    log_debug("Starting test_log_chat()")
     dump_hub_csv()
     log_chat(chat_id, session_id, model_name, prompt, history, info)
-    log_message("log_chat 1 returned")
+    log_debug("log_chat 1 returned")
     log_chat(chat_id, session_id, model_name, prompt + " + 2", history, info)
-    log_message("log_chat 2 returned")
+    log_debug("log_chat 2 returned")
     log_chat(chat_id, session_id, model_name, prompt + " + 3", history, info)
-    log_message("log_chat 3 returned")
+    log_debug("log_chat 3 returned")
     log_chat(chat_id, session_id, model_name, prompt + " + 4", history, info)
-    log_message("log_chat 4 returned")
+    log_debug("log_chat 4 returned")
 
     sleep_seconds = 10
-    log_message(f"Sleeping {sleep_seconds} seconds to let it finish and log the result.")
+    log_debug(f"Sleeping {sleep_seconds} seconds to let it finish and log the result.")
     time.sleep(sleep_seconds)
-    log_message("Finished sleeping.")
+    log_debug("Finished sleeping.")
     dump_hub_csv()
 
 

utils.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import time
 from functools import wraps
 from typing import Any, Literal
 
@@ -42,13 +43,28 @@ def get_model_config(model_name: str) -> dict:
     return config
 
 
-def log_message(message):
+def _log_message(prefix, message, icon=""):
+    timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
+    if len(icon) > 0:
+        icon = f"{icon} "
+    print(f"{timestamp}: {prefix} {icon}{message}")
+
+
+def log_debug(message):
     if DEBUG_MODE is True:
-        print(f"≫≫≫ {message}")
+        _log_message("DEBUG", message)
+
+
+def log_info(message):
+    _log_message("INFO ", message)
 
 
 def log_warning(message):
-    print(f"‼️ {message}")
+    _log_message("WARN ", message, "⚠️")
+
+
+def log_error(message):
+    _log_message("ERROR", message, "‼️")
 
 
 # Gradio 5.0.1 had issues with checking the message formats.  5.29.0 does not!
@@ -95,8 +111,8 @@ def logged_event_handler(log_msg='', event_handler=None, log_timer=None, clear_t
         if log_timer:
             if clear_timer:
                 log_timer.clear()
-            log_timer.add_step(f"Start: {log_message}")
-        log_message(f"::: Before event: {log_msg}")
+            log_timer.add_step(f"Start: {log_debug}")
+        log_debug(f"::: Before event: {log_msg}")
 
         # Call the original event handler
         result = event_handler(*args, **kwargs)
@@ -104,7 +120,7 @@ def logged_event_handler(log_msg='', event_handler=None, log_timer=None, clear_t
         # Log after
         if log_timer:
             log_timer.add_step(f"Completed: {log_msg}")
-        log_message(f"::: After event: {log_msg}")
+        log_debug(f"::: After event: {log_msg}")
 
         return result