diff --git a/modules/chat.py b/modules/chat.py index 881f7330..2db72f36 100644 --- a/modules/chat.py +++ b/modules/chat.py @@ -578,7 +578,6 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess visible_text = None stopping_strings = get_stopping_strings(state) is_stream = state['stream'] - # Prepare the input if not (regenerate or _continue): visible_text = html.escape(text) @@ -598,8 +597,6 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess # Apply extensions text, visible_text = apply_extensions('chat_input', text, visible_text, state) text = apply_extensions('input', text, state, is_chat=True) - - # Current row index output['internal'].append([text, '']) output['visible'].append([visible_text, '']) # Add metadata with timestamp @@ -644,16 +641,14 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess 'internal': output['internal'], 'metadata': output['metadata'] } - # Generate the prompt kwargs = { '_continue': _continue, '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) @@ -665,27 +660,53 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess # Generate reply = None for j, reply in enumerate(generate_reply(prompt, state, stopping_strings=stopping_strings, is_chat=True, for_ui=for_ui)): + # Handle start_with text (add to both internal and visible) + if state.get('start_with', '') and j == 0 and not _continue and not regenerate: + start_text = state['start_with'] + # Add to internal (English) version + reply = start_text + " " + reply + # Add to visible (translated) version + translated_start = apply_extensions('output', start_text, state, is_chat=True) + state['start_with'] = '' # Clear it after using # Extract the reply if state['mode'] in ['chat', 'chat-instruct']: visible_reply = re.sub("(||{{user}})", state['name1'], reply) else: visible_reply = reply - visible_reply = html.escape(visible_reply) - if shared.stop_everything: output['visible'][-1][1] = apply_extensions('output', output['visible'][-1][1], state, is_chat=True) yield output return if _continue: - output['internal'][-1] = [text, last_reply[0] + reply] - output['visible'][-1] = [visible_text, last_reply[1] + visible_reply] - elif not (j == 0 and visible_reply.strip() == ''): - output['internal'][-1] = [text, reply.lstrip(' ')] - output['visible'][-1] = [visible_text, visible_reply.lstrip(' ')] + # Separate already existing content from new content + original_internal = output['internal'][-1][1] + original_visible = output['visible'][-1][1] + # Get only the new generated part + new_content = reply[len(original_internal):] if reply.startswith(original_internal) else reply + new_content = new_content.lstrip() + + # Translate only the new part + translated_new = apply_extensions('output', new_content, state, is_chat=True) + + # Update both internal and visible versions + updated_internal = original_internal + " " + new_content + updated_visible = original_visible + " " + translated_new + + output['internal'][-1] = [text, updated_internal] + output['visible'][-1] = [visible_text, updated_visible] + + if is_stream: + yield output + elif not (j == 0 and visible_reply.strip() == ''): + # For normal generation, translate the whole reply + translated_reply = apply_extensions('output', visible_reply.lstrip(' '), state, is_chat=True) + output['internal'][-1] = [text, reply.lstrip(' ')] + output['visible'][-1] = [visible_text, translated_reply] + # Keep version metadata in sync during streaming (for regeneration) if regenerate: row_idx = len(output['internal']) - 1 @@ -695,11 +716,8 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess 'content': output['internal'][row_idx][1], 'visible_content': output['visible'][row_idx][1] }) - - if is_stream: - yield output - - output['visible'][-1][1] = apply_extensions('output', output['visible'][-1][1], state, is_chat=True) + if is_stream: + yield output # Final sync for version metadata (in case streaming was disabled) if regenerate: