Make ChatAssistant generic to support custom message types and add avatar provider support

src/main/java/com/flowingcode/vaadin/addons/chatassistant/ChatAssistant.java

@@ -21,7 +21,6 @@
 package com.flowingcode.vaadin.addons.chatassistant;
 
 import com.flowingcode.vaadin.addons.chatassistant.model.Message;
-import com.vaadin.flow.component.AttachEvent;
 import com.vaadin.flow.component.ClickNotifier;
 import com.vaadin.flow.component.Component;
 import com.vaadin.flow.component.ComponentEventListener;
@@ -43,6 +42,7 @@
 import com.vaadin.flow.data.provider.DataProvider;
 import com.vaadin.flow.data.renderer.ComponentRenderer;
 import com.vaadin.flow.data.renderer.Renderer;
+import com.vaadin.flow.function.SerializableSupplier;
 import com.vaadin.flow.shared.Registration;
 import java.time.LocalDateTime;
 import java.util.ArrayList;
@@ -64,21 +64,23 @@
 @JsModule("./react/animated-fab.tsx")
 @Tag("animated-fab")
 @CssImport("./styles/chat-assistant-styles.css")
-public class ChatAssistant extends ReactAdapterComponent implements ClickNotifier<ChatAssistant> {
+public class ChatAssistant<T extends Message> extends ReactAdapterComponent implements ClickNotifier<ChatAssistant<T>> {
 
   private static final String CHAT_HEADER_CLASS_NAME = "chat-header";
   private static final String PADDING_SMALL = "0.5em";
 
   private Component headerComponent;
   private VerticalLayout container;
   private Component footerContainer;
-  private VirtualList<Message> content = new VirtualList<>();
+  private VirtualList<T> content = new VirtualList<>();
   private Popover chatWindow;
-  private List<Message> messages;
+  private List<T> messages;
   private MessageInput messageInput;
   private Span whoIsTyping;
   private boolean minimized = false;
   private Registration defaultSubmitListenerRegistration;
+  private SerializableSupplier<Avatar> avatarProvider = () -> new Avatar("Chat Assistant");
+  private Avatar avatar;
 
   /**
    * Default constructor. Creates a ChatAssistant with no messages.
@@ -102,7 +104,7 @@ public ChatAssistant(boolean markdownEnabled) {
    * @param messages the list of messages
    * @param markdownEnabled flag to enable or disable markdown support
    */
-  public ChatAssistant(List<Message> messages, boolean markdownEnabled) {
+  public ChatAssistant(List<T> messages, boolean markdownEnabled) {
     this.messages = messages;
     initializeHeader();
     initializeFooter();
@@ -121,13 +123,15 @@ private void initializeHeader() {
     headerComponent = header;
   }
 
+  @SuppressWarnings("unchecked")
   private void initializeFooter() {
     messageInput = new MessageInput();
     messageInput.setSizeFull();
     messageInput.getStyle().set("padding", PADDING_SMALL);
-    defaultSubmitListenerRegistration = messageInput.addSubmitListener(se -> 
-        sendMessage(Message.builder().messageTime(LocalDateTime.now())
-            .name("User").content(se.getValue()).build()));
+    defaultSubmitListenerRegistration = messageInput.addSubmitListener(se -> {
+      sendMessage((T) Message.builder().messageTime(LocalDateTime.now())
+          .name("User").content(se.getValue()).build());
+    });
     whoIsTyping = new Span();
     whoIsTyping.setClassName("chat-assistant-who-is-typing");
     whoIsTyping.setVisible(false);
@@ -139,13 +143,14 @@ private void initializeFooter() {
     footerContainer = footer;
   }
 
+  @SuppressWarnings("unchecked")
   private void initializeContent(boolean markdownEnabled) {
-    content.setItems(messages);
-    content.setRenderer(new ComponentRenderer<>(message -> new ChatMessage(message, markdownEnabled), 
+    content.setRenderer(new ComponentRenderer<>(message -> new ChatMessage<T>(message, markdownEnabled), 
         (component, message) -> {
-          ((ChatMessage) component).setMessage(message);
+          ((ChatMessage<T>) component).setMessage(message);
           return component;
         }));
+    content.setItems(messages);
     content.setMinHeight("400px");
     content.setMinWidth("400px");
     container = new VerticalLayout(headerComponent, content, footerContainer);
@@ -177,7 +182,10 @@ private void initializeChatWindow() {
   }
 
   private void initializeAvatar() {
-    Avatar avatar = new Avatar("AI");
+    if (avatar!=null) {
+      avatar.removeFromParent();
+    }
+    avatar = avatarProvider.get();
     avatar.setSizeFull();
     this.getElement().appendChild(avatar.getElement());
     this.addAttachListener(ev -> this.getElement().executeJs("return;")
@@ -190,7 +198,7 @@ private void initializeAvatar() {
    * 
    * @param dataProvider the data provider to be used
    */
-  public void setDataProvider(DataProvider<Message, ?> dataProvider) {
+  public void setDataProvider(DataProvider<T, ?> dataProvider) {
     content.setDataProvider(dataProvider);
   }
 
@@ -243,7 +251,7 @@ private void refreshContent() {
    *
    * @param message the message to be sent programmatically
    */
-  public void sendMessage(Message message) {
+  public void sendMessage(T message) {
     messages.add(message);
     content.getDataProvider().refreshAll();
     content.scrollToEnd();
@@ -254,7 +262,7 @@ public void sendMessage(Message message) {
    * 
    * @param message the message to be updated
    */
-  public void updateMessage(Message message) {
+  public void updateMessage(T message) {
     this.content.getDataProvider().refreshItem(message);
   }
 
@@ -363,9 +371,19 @@ public void scrollToEnd() {
    *
    * @param renderer the renderer to use for rendering {@link Message} objects, it cannot be null
    */
-  public void setMessagesRenderer(Renderer<Message> renderer) {
+  public void setMessagesRenderer(Renderer<T> renderer) {
     Objects.requireNonNull(renderer, "Renderer cannot not be null");
     content.setRenderer(renderer);
   }
 
+  /**
+   * Sets the avatar provider that will be used to create the avatar
+   * 
+   * @param avatarProvider
+   */
+  public void setAvatarProvider(SerializableSupplier<Avatar> avatarProvider) {
+    this.avatarProvider = avatarProvider;
+    this.initializeAvatar();
+  }
+
 }

src/main/java/com/flowingcode/vaadin/addons/chatassistant/ChatMessage.java

@@ -41,9 +41,9 @@
 @Tag("vaadin-message")
 @CssImport("./styles/chat-message-styles.css")
 @EqualsAndHashCode(callSuper=false)
-public class ChatMessage extends Component implements HasComponents {
+public class ChatMessage<T extends Message> extends Component implements HasComponents {
 
-  private Message message;
+  private T message;
   private boolean markdownEnabled;
   private Div loader;
 
@@ -52,7 +52,7 @@ public class ChatMessage extends Component implements HasComponents {
    * 
    * @param message message used to populate the ChatMessage instance
    */
-  public ChatMessage(Message message) {
+  public ChatMessage(T message) {
     this(message, false);
   }
 
@@ -62,7 +62,7 @@ public ChatMessage(Message message) {
    * @param message message used to populate the ChatMessage instance
    * @param markdownEnabled whether the message supports markdown or not
    */
-  public ChatMessage(Message message, boolean markdownEnabled) {
+  public ChatMessage(T message, boolean markdownEnabled) {
     this.markdownEnabled = markdownEnabled;
     setMessage(message);
   }
@@ -72,7 +72,7 @@ public ChatMessage(Message message, boolean markdownEnabled) {
    * 
    * @param message message used to populate the ChatMessage instance
    */
-  public void setMessage(Message message) {
+  public void setMessage(T message) {
     this.message = message;
     updateLoadingState(message);
     if (message.getName()!=null) {
@@ -87,7 +87,7 @@ public void setMessage(Message message) {
     }
   }
 
-  private void updateLoadingState(Message message) {
+  private void updateLoadingState(T message) {
     if (message.isLoading()) {
       loader = new Div(new Div(),new Div(), new Div(), new Div());
       loader.setClassName("lds-ellipsis");
@@ -112,7 +112,7 @@ private void updateLoadingState(Message message) {
    * 
    * @return the message object used to populate this ChatMessage
    */
-  public Message getMessage() {
+  public T getMessage() {
     return message;
   }
 

src/main/java/com/flowingcode/vaadin/addons/chatassistant/model/Message.java

@@ -26,6 +26,7 @@
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.Setter;
+import lombok.experimental.SuperBuilder;
 
 /**
  * Class that represents a chat message
@@ -35,7 +36,7 @@
 @SuppressWarnings("serial")
 @Getter
 @Setter
-@Builder
+@SuperBuilder
 @EqualsAndHashCode(of = "id")
 public class Message implements Serializable {
 

src/test/java/com/flowingcode/vaadin/addons/chatassistant/ChatAssistantDemo.java

@@ -19,15 +19,16 @@
  */
 package com.flowingcode.vaadin.addons.chatassistant;
 
-import com.flowingcode.vaadin.addons.chatassistant.model.Message;
 import com.flowingcode.vaadin.addons.demo.DemoSource;
 import com.flowingcode.vaadin.addons.demo.SourcePosition;
 import com.google.common.base.Strings;
 import com.vaadin.flow.component.UI;
+import com.vaadin.flow.component.avatar.Avatar;
 import com.vaadin.flow.component.button.Button;
 import com.vaadin.flow.component.dependency.CssImport;
 import com.vaadin.flow.component.orderedlayout.VerticalLayout;
 import com.vaadin.flow.component.textfield.TextArea;
+import com.vaadin.flow.data.renderer.ComponentRenderer;
 import com.vaadin.flow.router.PageTitle;
 import com.vaadin.flow.router.Route;
 import java.time.LocalDateTime;
@@ -41,7 +42,8 @@
 public class ChatAssistantDemo extends VerticalLayout {
 
   public ChatAssistantDemo() {
-    ChatAssistant chatAssistant = new ChatAssistant();
+    ChatAssistant<CustomMessage> chatAssistant = new ChatAssistant<>();
+    chatAssistant.setAvatarProvider(()->new Avatar("Chat Assistant","chatbot.png"));
     TextArea message = new TextArea();
     message.setLabel("Enter a message from the assistant");
     message.setSizeFull();
@@ -51,20 +53,31 @@ public ChatAssistantDemo() {
       }
     });
     message.addBlurListener(ev->chatAssistant.clearWhoIsTyping());
+    chatAssistant.setMessagesRenderer(new ComponentRenderer<CustomChatMessage,CustomMessage>(m -> {
+      return new CustomChatMessage(m); 
+    },
+        (component, m) -> {
+          ((CustomChatMessage) component).setMessage(m);
+          return component;
+        }));
+    chatAssistant.setSubmitListener(se -> {
+      chatAssistant.sendMessage(CustomMessage.builder().messageTime(LocalDateTime.now())
+          .name("User").content(se.getValue()).tagline("Generated by user").build());
+    });
 
     Button chat = new Button("Chat");
     chat.addClickListener(ev -> {
-      Message m = Message.builder().content(message.getValue()).messageTime(LocalDateTime.now())
-          .name("Assistant").avatar("chatbot.png").build();
+      CustomMessage m = CustomMessage.builder().content(message.getValue()).messageTime(LocalDateTime.now())
+          .name("Assistant").avatar("chatbot.png").tagline("Generated by assistant").build();
 
       chatAssistant.sendMessage(m);
       message.clear();
     });
     Button chatWithThinking = new Button("Chat With Thinking");
     chatWithThinking.addClickListener(ev -> {
-      Message delayedMessage = Message.builder().loading(true).content(message.getValue())
+      CustomMessage delayedMessage = CustomMessage.builder().loading(true).content(message.getValue())
           .messageTime(LocalDateTime.now())
-          .name("Assistant").avatar("chatbot.png").build();
+          .name("Assistant").avatar("chatbot.png").tagline("Generated by assistant").build();
 
       UI currentUI = UI.getCurrent();
       chatAssistant.sendMessage(delayedMessage);
@@ -81,9 +94,9 @@ public void run() {
 
       message.clear();
     });
-    chatAssistant.sendMessage(Message.builder().content("Hello, I am here to assist you")
+    chatAssistant.sendMessage(CustomMessage.builder().content("Hello, I am here to assist you")
         .messageTime(LocalDateTime.now())
-        .name("Assistant").avatar("chatbot.png").build());
+        .name("Assistant").avatar("chatbot.png").tagline("Generated by assistant").build());
 
     add(message, chat, chatWithThinking, chatAssistant);
   }

src/test/java/com/flowingcode/vaadin/addons/chatassistant/ChatAssistantLazyLoadingDemo.java

@@ -123,7 +123,7 @@ public class ChatAssistantLazyLoadingDemo extends VerticalLayout {
       ));
 
   public ChatAssistantLazyLoadingDemo() {
-    ChatAssistant chatAssistant = new ChatAssistant();
+    ChatAssistant<Message> chatAssistant = new ChatAssistant<>();
     chatAssistant.setClassName("small");
     Span lazyLoadingData = new Span();
     DataProvider<Message,?> dataProvider = DataProvider.fromCallbacks(query->{

src/test/java/com/flowingcode/vaadin/addons/chatassistant/ChatAssistantMarkdownDemo.java

@@ -23,19 +23,13 @@
 import com.flowingcode.vaadin.addons.demo.DemoSource;
 import com.flowingcode.vaadin.addons.demo.SourcePosition;
 import com.google.common.base.Strings;
-import com.vaadin.flow.component.UI;
 import com.vaadin.flow.component.button.Button;
 import com.vaadin.flow.component.dependency.CssImport;
-import com.vaadin.flow.component.html.Span;
-import com.vaadin.flow.component.icon.Icon;
-import com.vaadin.flow.component.icon.VaadinIcon;
-import com.vaadin.flow.component.orderedlayout.HorizontalLayout;
 import com.vaadin.flow.component.orderedlayout.VerticalLayout;
 import com.vaadin.flow.component.textfield.TextArea;
 import com.vaadin.flow.router.PageTitle;
 import com.vaadin.flow.router.Route;
 import java.time.LocalDateTime;
-import java.util.Timer;
 
 @DemoSource(sourcePosition = SourcePosition.PRIMARY)
 @PageTitle("Markdown Demo")
@@ -45,7 +39,7 @@
 public class ChatAssistantMarkdownDemo extends VerticalLayout {
 
   public ChatAssistantMarkdownDemo() {
-    ChatAssistant chatAssistant = new ChatAssistant(true);
+    ChatAssistant<Message> chatAssistant = new ChatAssistant<>(true);
     TextArea message = new TextArea();
     message.setLabel("Enter a message from the assistant (try using Markdown)");
     message.setSizeFull();

src/test/java/com/flowingcode/vaadin/addons/chatassistant/CustomChatMessage.java

@@ -0,0 +1,20 @@
+package com.flowingcode.vaadin.addons.chatassistant;
+
+import com.vaadin.flow.component.html.Span;
+
+@SuppressWarnings("serial")
+public class CustomChatMessage extends ChatMessage<CustomMessage> {
+
+  private Span tagline = new Span();
+
+  public CustomChatMessage(CustomMessage message) {
+    super(message);
+    tagline.setText(message.getTagline());
+    tagline.getStyle().set("display", "block");
+    tagline.getStyle().set("font-size", "x-small");
+    this.add(tagline);
+  }
+
+
+
+}

src/test/java/com/flowingcode/vaadin/addons/chatassistant/CustomMessage.java

@@ -0,0 +1,18 @@
+package com.flowingcode.vaadin.addons.chatassistant;
+
+import com.flowingcode.vaadin.addons.chatassistant.model.Message;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.experimental.SuperBuilder;
+
+@SuppressWarnings("serial")
+@EqualsAndHashCode(callSuper = true)
+@Getter
+@Setter
+@SuperBuilder
+public class CustomMessage extends Message {
+
+  private String tagline;
+
+}