From eb36492ba55f82a18cd7c8baaba4f61c324a6e49 Mon Sep 17 00:00:00 2001 From: sabaimran Date: Sun, 8 Dec 2024 21:46:04 -0800 Subject: [PATCH] Update handling of images when included in the chat history with assistant message --- src/khoj/processor/conversation/utils.py | 27 ++++++++++++------------ 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/src/khoj/processor/conversation/utils.py b/src/khoj/processor/conversation/utils.py index d68e592e..12fc736f 100644 --- a/src/khoj/processor/conversation/utils.py +++ b/src/khoj/processor/conversation/utils.py @@ -433,20 +433,19 @@ def generate_chatml_messages_with_context( reconstructed_context_message = ChatMessage(content=message_context, role="user") chatml_messages.insert(0, reconstructed_context_message) - if chat.get("images"): - if role == "assistant": - # Issue: the assistant role cannot accept an image as a message content, so send it in a separate user message. - file_attachment_message = construct_structured_message( - message=prompts.generated_image_attachment.format(), - images=chat.get("images"), - model_type=model_type, - vision_enabled=vision_enabled, - ) - chatml_messages.append(ChatMessage(content=file_attachment_message, role="user")) - else: - message_content = construct_structured_message( - chat_message, chat.get("images"), model_type, vision_enabled - ) + if chat.get("images") and role == "assistant": + # Issue: the assistant role cannot accept an image as a message content, so send it in a separate user message. + file_attachment_message = construct_structured_message( + message=prompts.generated_image_attachment.format(), + images=chat.get("images"), + model_type=model_type, + vision_enabled=vision_enabled, + ) + chatml_messages.append(ChatMessage(content=file_attachment_message, role="user")) + + message_content = construct_structured_message( + chat_message, chat.get("images") if role == "user" else [], model_type, vision_enabled + ) reconstructed_message = ChatMessage(content=message_content, role=role) chatml_messages.insert(0, reconstructed_message)