From 8e10f9894aee28088b176adc5bed91f80cc11b5c Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 20 May 2025 10:48:46 -0300 Subject: [PATCH] Add a metadata field to the chat history & add date/time to chat messages (#7003) --- css/main.css | 33 +++++++++++ modules/chat.py | 114 ++++++++++++++++++++++++++++++++++++-- modules/html_generator.py | 65 +++++++++++++++++++--- modules/ui_chat.py | 2 +- 4 files changed, 200 insertions(+), 14 deletions(-) diff --git a/css/main.css b/css/main.css index dc0ba817..319c1778 100644 --- a/css/main.css +++ b/css/main.css @@ -419,6 +419,14 @@ div.svelte-362y77>*, div.svelte-362y77>.form>* { padding-right: 1rem; } +.chat .message .timestamp { + font-size: 0.7em; + display: inline-block; + font-weight: normal; + opacity: 0.7; + margin-left: 5px; +} + .chat-parent.bigchat { flex: 1; } @@ -1269,6 +1277,31 @@ div.svelte-362y77>*, div.svelte-362y77>.form>* { left: 75px; } +.footer-button.footer-info-button { + bottom: -23px; +} + +.user-message .footer-button.footer-info-button { + left: 25px; +} + +.assistant-message:not(:last-child) .footer-button.footer-info-button { + left: 25px; +} + +.assistant-message:last-child .footer-button.footer-info-button { + left: 100px; +} + +.message:not(:last-child) .text-bot .footer-button.footer-info-button, +.message .text-you .footer-button.footer-info-button { + left: 25px; +} + +.message:last-child .text-bot .footer-button.footer-info-button { + left: 100px; +} + .message:hover .footer-button, .user-message:hover .footer-button, .assistant-message:hover .footer-button { diff --git a/modules/chat.py b/modules/chat.py index b83c4bfe..cbcde212 100644 --- a/modules/chat.py +++ b/modules/chat.py @@ -37,6 +37,30 @@ def strftime_now(format): return datetime.now().strftime(format) +def get_current_timestamp(): + """Returns the current time in 24-hour format""" + return datetime.now().strftime('%b %d, %Y %H:%M') + + +def update_message_metadata(metadata_dict, role, index, **fields): + """ + Updates or adds metadata fields for a specific message. + + Args: + metadata_dict: The metadata dictionary + role: The role (user, assistant, etc) + index: The message index + **fields: Arbitrary metadata fields to update/add + """ + key = f"{role}_{index}" + if key not in metadata_dict: + metadata_dict[key] = {} + + # Update with provided fields + for field_name, field_value in fields.items(): + metadata_dict[key][field_name] = field_value + + jinja_env = ImmutableSandboxedEnvironment( trim_blocks=True, lstrip_blocks=True, @@ -347,6 +371,10 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess output = apply_extensions('history', output) state = apply_extensions('state', state) + # Initialize metadata if not present + if 'metadata' not in output: + output['metadata'] = {} + visible_text = None stopping_strings = get_stopping_strings(state) is_stream = state['stream'] @@ -359,40 +387,56 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess text, visible_text = apply_extensions('chat_input', text, visible_text, state) text = apply_extensions('input', text, state, is_chat=True) + # Current row index + row_idx = len(output['internal']) output['internal'].append([text, '']) output['visible'].append([visible_text, '']) + # Add metadata with timestamp + update_message_metadata(output['metadata'], "user", row_idx, timestamp=get_current_timestamp()) # *Is typing...* if loading_message: yield { 'visible': output['visible'][:-1] + [[output['visible'][-1][0], shared.processing_message]], - 'internal': output['internal'] + 'internal': output['internal'], + 'metadata': output['metadata'] } else: text, visible_text = output['internal'][-1][0], output['visible'][-1][0] if regenerate: + row_idx = len(output['internal']) - 1 if loading_message: yield { 'visible': output['visible'][:-1] + [[visible_text, shared.processing_message]], - 'internal': output['internal'][:-1] + [[text, '']] + 'internal': output['internal'][:-1] + [[text, '']], + 'metadata': output['metadata'] } elif _continue: last_reply = [output['internal'][-1][1], output['visible'][-1][1]] if loading_message: yield { 'visible': output['visible'][:-1] + [[visible_text, last_reply[1] + '...']], - 'internal': output['internal'] + 'internal': output['internal'], + 'metadata': output['metadata'] } # Generate the prompt kwargs = { '_continue': _continue, - 'history': output if _continue else {k: v[:-1] for k, v in output.items()} + 'history': output if _continue else { + k: (v[:-1] if k in ['internal', 'visible'] else v) + for k, v in output.items() + } } + prompt = apply_extensions('custom_generate_chat_prompt', text, state, **kwargs) if prompt is None: prompt = generate_chat_prompt(text, state, **kwargs) + # Add timestamp for assistant's response at the start of generation + row_idx = len(output['internal']) - 1 + update_message_metadata(output['metadata'], "assistant", row_idx, timestamp=get_current_timestamp()) + # Generate reply = None for j, reply in enumerate(generate_reply(prompt, state, stopping_strings=stopping_strings, is_chat=True, for_ui=for_ui)): @@ -495,9 +539,19 @@ def generate_chat_reply_wrapper(text, state, regenerate=False, _continue=False): def remove_last_message(history): + if 'metadata' not in history: + history['metadata'] = {} + if len(history['visible']) > 0 and history['internal'][-1][0] != '<|BEGIN-VISIBLE-CHAT|>': + row_idx = len(history['internal']) - 1 last = history['visible'].pop() history['internal'].pop() + + # Remove metadata directly by known keys + if f"user_{row_idx}" in history['metadata']: + del history['metadata'][f"user_{row_idx}"] + if f"assistant_{row_idx}" in history['metadata']: + del history['metadata'][f"assistant_{row_idx}"] else: last = ['', ''] @@ -514,30 +568,54 @@ def send_last_reply_to_input(history): def replace_last_reply(text, state): history = state['history'] + # Initialize metadata if not present + if 'metadata' not in history: + history['metadata'] = {} + if len(text.strip()) == 0: return history elif len(history['visible']) > 0: + row_idx = len(history['internal']) - 1 history['visible'][-1][1] = html.escape(text) history['internal'][-1][1] = apply_extensions('input', text, state, is_chat=True) + update_message_metadata(history['metadata'], "assistant", row_idx, timestamp=get_current_timestamp()) return history def send_dummy_message(text, state): history = state['history'] + + # Initialize metadata if not present + if 'metadata' not in history: + history['metadata'] = {} + + row_idx = len(history['internal']) history['visible'].append([html.escape(text), '']) history['internal'].append([apply_extensions('input', text, state, is_chat=True), '']) + update_message_metadata(history['metadata'], "user", row_idx, timestamp=get_current_timestamp()) + return history def send_dummy_reply(text, state): history = state['history'] + + # Initialize metadata if not present + if 'metadata' not in history: + history['metadata'] = {} + if len(history['visible']) > 0 and not history['visible'][-1][1] == '': + row_idx = len(history['internal']) history['visible'].append(['', '']) history['internal'].append(['', '']) + # We don't need to add system metadata + row_idx = len(history['internal']) - 1 history['visible'][-1][1] = html.escape(text) history['internal'][-1][1] = apply_extensions('input', text, state, is_chat=True) + update_message_metadata(history['metadata'], "assistant", row_idx, timestamp=get_current_timestamp()) + return history @@ -547,7 +625,8 @@ def redraw_html(history, name1, name2, mode, style, character, reset_cache=False def start_new_chat(state): mode = state['mode'] - history = {'internal': [], 'visible': []} + # Initialize with empty metadata dictionary + history = {'internal': [], 'visible': [], 'metadata': {}} if mode != 'instruct': greeting = replace_character_names(state['greeting'], state['name1'], state['name2']) @@ -555,6 +634,9 @@ def start_new_chat(state): history['internal'] += [['<|BEGIN-VISIBLE-CHAT|>', greeting]] history['visible'] += [['', apply_extensions('output', html.escape(greeting), state, is_chat=True)]] + # Add timestamp for assistant's greeting + update_message_metadata(history['metadata'], "assistant", 0, timestamp=get_current_timestamp()) + unique_id = datetime.now().strftime('%Y%m%d-%H-%M-%S') save_history(history, unique_id, state['character_menu'], state['mode']) @@ -735,6 +817,16 @@ def load_history(unique_id, character, mode): 'visible': f['data_visible'] } + # Add metadata if it doesn't exist + if 'metadata' not in history: + history['metadata'] = {} + # Add placeholder timestamps for existing messages + for i, (user_msg, asst_msg) in enumerate(history['internal']): + if user_msg and user_msg != '<|BEGIN-VISIBLE-CHAT|>': + update_message_metadata(history['metadata'], "user", i, timestamp="") + if asst_msg: + update_message_metadata(history['metadata'], "assistant", i, timestamp="") + return history @@ -750,6 +842,16 @@ def load_history_json(file, history): 'visible': f['data_visible'] } + # Add metadata if it doesn't exist + if 'metadata' not in history: + history['metadata'] = {} + # Add placeholder timestamps + for i, (user_msg, asst_msg) in enumerate(history['internal']): + if user_msg and user_msg != '<|BEGIN-VISIBLE-CHAT|>': + update_message_metadata(history['metadata'], "user", i, timestamp="") + if asst_msg: + update_message_metadata(history['metadata'], "assistant", i, timestamp="") + return history except: return history @@ -1299,7 +1401,7 @@ def handle_your_picture_change(picture, state): def handle_send_instruction_click(state): state['mode'] = 'instruct' - state['history'] = {'internal': [], 'visible': []} + state['history'] = {'internal': [], 'visible': [], 'metadata': {}} output = generate_chat_prompt("Input", state) diff --git a/modules/html_generator.py b/modules/html_generator.py index 39659476..5dbde6da 100644 --- a/modules/html_generator.py +++ b/modules/html_generator.py @@ -169,11 +169,7 @@ def convert_to_markdown(string, message_id=None): thinking_block = f'''
- - - - - + {info_svg_small} {title_text}
{thinking_html}
@@ -339,11 +335,24 @@ copy_svg = '''''' continue_svg = '''''' remove_svg = '''''' +info_svg = '''''' +info_svg_small = '''''' copy_button = f'' refresh_button = f'' continue_button = f'' remove_button = f'' +info_button = f'' + + +def format_message_timestamp(history, role, index): + """Get a formatted timestamp HTML span for a message if available""" + key = f"{role}_{index}" + if 'metadata' in history and key in history['metadata'] and history['metadata'][key].get('timestamp'): + timestamp = history['metadata'][key]['timestamp'] + return f"{timestamp}" + + return "" def generate_instruct_html(history): @@ -354,6 +363,23 @@ def generate_instruct_html(history): row_internal = history['internal'][i] converted_visible = [convert_to_markdown_wrapped(entry, message_id=i, use_cache=i != len(history['visible']) - 1) for entry in row_visible] + # Get timestamps + user_timestamp = format_message_timestamp(history, "user", i) + assistant_timestamp = format_message_timestamp(history, "assistant", i) + + # Create info buttons for timestamps if they exist + info_message_user = "" + if user_timestamp != "": + # Extract the timestamp value from the span + user_timestamp_value = user_timestamp.split('>', 1)[1].split('<', 1)[0] + info_message_user = info_button.replace("message", user_timestamp_value) + + info_message_assistant = "" + if assistant_timestamp != "": + # Extract the timestamp value from the span + assistant_timestamp_value = assistant_timestamp.split('>', 1)[1].split('<', 1)[0] + info_message_assistant = info_button.replace("message", assistant_timestamp_value) + if converted_visible[0]: # Don't display empty user messages output += ( f'
' f'
{converted_visible[0]}
' f'{copy_button}' + f'{info_message_user}' f'
' f'' ) @@ -374,6 +401,7 @@ def generate_instruct_html(history): f'{refresh_button if i == len(history["visible"]) - 1 else ""}' f'{continue_button if i == len(history["visible"]) - 1 else ""}' f'{remove_button if i == len(history["visible"]) - 1 else ""}' + f'{info_message_assistant}' f'' f'' ) @@ -401,13 +429,17 @@ def generate_cai_chat_html(history, name1, name2, style, character, reset_cache= row_internal = history['internal'][i] converted_visible = [convert_to_markdown_wrapped(entry, message_id=i, use_cache=i != len(history['visible']) - 1) for entry in row_visible] + # Get timestamps + user_timestamp = format_message_timestamp(history, "user", i) + assistant_timestamp = format_message_timestamp(history, "assistant", i) + if converted_visible[0]: # Don't display empty user messages output += ( f'
' f'
{img_me}
' f'
' - f'
{name1}
' + f'
{name1}{user_timestamp}
' f'
{converted_visible[0]}
' f'{copy_button}' f'
' @@ -419,7 +451,7 @@ def generate_cai_chat_html(history, name1, name2, style, character, reset_cache= f'data-raw="{html.escape(row_internal[1], quote=True)}">' f'
{img_bot}
' f'
' - f'
{name2}
' + f'
{name2}{assistant_timestamp}
' f'
{converted_visible[1]}
' f'{copy_button}' f'{refresh_button if i == len(history["visible"]) - 1 else ""}' @@ -441,6 +473,23 @@ def generate_chat_html(history, name1, name2, reset_cache=False): row_internal = history['internal'][i] converted_visible = [convert_to_markdown_wrapped(entry, message_id=i, use_cache=i != len(history['visible']) - 1) for entry in row_visible] + # Get timestamps + user_timestamp = format_message_timestamp(history, "user", i) + assistant_timestamp = format_message_timestamp(history, "assistant", i) + + # Create info buttons for timestamps if they exist + info_message_user = "" + if user_timestamp != "": + # Extract the timestamp value from the span + user_timestamp_value = user_timestamp.split('>', 1)[1].split('<', 1)[0] + info_message_user = info_button.replace("message", user_timestamp_value) + + info_message_assistant = "" + if assistant_timestamp != "": + # Extract the timestamp value from the span + assistant_timestamp_value = assistant_timestamp.split('>', 1)[1].split('<', 1)[0] + info_message_assistant = info_button.replace("message", assistant_timestamp_value) + if converted_visible[0]: # Don't display empty user messages output += ( f'
' f'
{converted_visible[0]}
' f'{copy_button}' + f'{info_message_user}' f'
' f'
' ) @@ -461,6 +511,7 @@ def generate_chat_html(history, name1, name2, reset_cache=False): f'{refresh_button if i == len(history["visible"]) - 1 else ""}' f'{continue_button if i == len(history["visible"]) - 1 else ""}' f'{remove_button if i == len(history["visible"]) - 1 else ""}' + f'{info_message_assistant}' f'
' f'' ) diff --git a/modules/ui_chat.py b/modules/ui_chat.py index a0c37dad..7a5430ca 100644 --- a/modules/ui_chat.py +++ b/modules/ui_chat.py @@ -47,7 +47,7 @@ def create_ui(): with gr.Row(): with gr.Column(elem_id='chat-col'): shared.gradio['display'] = gr.JSON(value={}, visible=False) # Hidden buffer - shared.gradio['html_display'] = gr.HTML(value=chat_html_wrapper({'internal': [], 'visible': []}, '', '', 'chat', 'cai-chat', '')['html'], visible=True) + shared.gradio['html_display'] = gr.HTML(value=chat_html_wrapper({'internal': [], 'visible': [], 'metadata': {}}, '', '', 'chat', 'cai-chat', '')['html'], visible=True) with gr.Row(elem_id="chat-input-row"): with gr.Column(scale=1, elem_id='gr-hover-container'): gr.HTML(value='
☰
', elem_id='gr-hover')