better styling, now thoughts and actions are hidden, scroll works better

next steps:

- [ ] Ensure it doesn’t create more web documents when one already exists
- [ ] Citations should not be rendered on the next line but on the same line as the text
- [ ] If invalid XML, run get 3.5 to verify and fix XML based one examples
- [ ] Making sure if you ask for other information, it doesn’t go to the same website. Providing website history in use rules for the search tool and website scraper tool or in the prompt directly

6 files changed, 352 insertions, 324 deletions

diff --git a/src/client/util/CurrentUserUtils.ts b/src/client/util/CurrentUserUtils.ts
index e095bc659..280830442 100644
--- a/src/client/util/CurrentUserUtils.ts
+++ b/src/client/util/CurrentUserUtils.ts
@@ -371,7 +371,7 @@ pie title Minerals in my tap water
             {key: "Button",     creator: Docs.Create.ButtonDocument,                        opts: { _width: 150, _height: 50, _xPadding: 10, _yPadding: 10, title_custom: true,  waitForDoubleClickToClick: 'never'}, scripts: {onClick: FollowLinkScript()?.script.originalScript ?? ""}},
             {key: "Script",     creator: opts => Docs.Create.ScriptingDocument(null, opts), opts: { _width: 200, _height: 250, }},
             {key: "DataViz",    creator: opts => Docs.Create.DataVizDocument("/users/rz/Downloads/addresses.csv", opts),   opts: { _width: 300, _height: 300 }},
-            {key: "Chat",       creator: Docs.Create.ChatDocument,          opts: { _width: 300, _height: 300, }},
+            {key: "Chat",       creator: Docs.Create.ChatDocument,          opts: { _width: 500, _height: 500, }},
             {key: "Header",     creator: headerTemplate,                                    opts: { _width: 300, _height: 120, _header_pointerEvents: "all", _header_height: 50, _header_fontSize: 9,_layout_autoHeightMargins: 50, _layout_autoHeight: true, treeView_HideUnrendered: true}},
             {key: "ViewSlide",  creator: slideView,                                         opts: { _width: 400, _height: 300, _xMargin: 3, _yMargin: 3,}},
             {key: "Trail",      creator: Docs.Create.PresDocument,                          opts: { _width: 400, _height: 30, _type_collection: CollectionViewType.Stacking, dropAction: dropActionType.embed, treeView_HideTitle: true, _layout_fitWidth:true, layout_boxShadow: "0 0" }},
diff --git a/src/client/views/nodes/ChatBox/ChatBox.scss b/src/client/views/nodes/ChatBox/ChatBox.scss
index 91bb3aba7..76fa05ce8 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.scss
+++ b/src/client/views/nodes/ChatBox/ChatBox.scss
@@ -1,246 +1,246 @@
-$background-color: #f8f9fa;
+@import url('https://fonts.googleapis.com/css2?family=Atkinson+Hyperlegible:ital,wght@0,400;0,700;1,400;1,700&display=swap');
+
+$primary-color: #4a90e2;
+$secondary-color: #f5f8fa;
 $text-color: #333;
-$input-background: #fff;
-$button-color: #007bff;
-$button-hover-color: darken($button-color, 10%);
-$shadow-color: rgba(0, 0, 0, 0.075);
-$border-radius: 8px;
-$citation-color: #ff6347;
-$citation-hover-color: darken($citation-color, 10%);
-$follow-up-bg-color: #e9ecef;
-$follow-up-hover-bg-color: #dee2e6;
-
-.chatBox {
+$light-text-color: #777;
+$border-color: #e1e8ed;
+$shadow-color: rgba(0, 0, 0, 0.1);
+$transition: all 0.3s ease;
+
+.chat-box {
     display: flex;
     flex-direction: column;
-    width: 100%;
     height: 100%;
-    background-color: $background-color;
-    font-family: 'Helvetica Neue', Arial, sans-serif;
+    background-color: #fff;
+    font-family:
+        'Atkinson Hyperlegible',
+        -apple-system,
+        BlinkMacSystemFont,
+        'Segoe UI',
+        Roboto,
+        Helvetica,
+        Arial,
+        sans-serif;
+    border-radius: 12px;
+    overflow: hidden;
+    box-shadow: 0 4px 12px $shadow-color;
+
+    .chat-header {
+        background-color: $primary-color;
+        color: white;
+        padding: 15px;
+        text-align: center;
+        box-shadow: 0 2px 4px $shadow-color;
+        height: fit-content;
+        h2 {
+            margin: 0;
+            font-size: 1.3em;
+            font-weight: 500;
+        }
+    }
 
-    .scroll-box {
+    .chat-messages {
         flex-grow: 1;
-        overflow-y: scroll;
-        overflow-x: hidden;
-        height: 100%;
-        padding: 10px;
+        overflow-y: auto;
+        padding: 20px;
         display: flex;
-        flex-direction: column-reverse;
-        padding-bottom: 0;
+        flex-direction: column;
 
         &::-webkit-scrollbar {
-            width: 8px;
+            width: 6px;
         }
+
         &::-webkit-scrollbar-thumb {
-            background-color: darken($background-color, 10%);
-            border-radius: $border-radius;
+            background-color: $border-color;
+            border-radius: 3px;
         }
+    }
 
-        .chat-content {
-            display: flex;
-            flex-direction: column;
+    .chat-input {
+        display: flex;
+        padding: 20px;
+        border-top: 1px solid $border-color;
+        background-color: #fff;
+
+        input {
+            flex-grow: 1;
+            padding: 12px 15px;
+            border: 1px solid $border-color;
+            border-radius: 24px;
+            font-size: 15px;
+            transition: $transition;
+
+            &:focus {
+                outline: none;
+                border-color: $primary-color;
+                box-shadow: 0 0 0 2px rgba($primary-color, 0.2);
+            }
         }
 
-        .messages {
+        .submit-button {
+            background-color: $primary-color;
+            color: white;
+            border: none;
+            border-radius: 50%;
+            width: 48px;
+            height: 48px;
+            margin-left: 10px;
+            cursor: pointer;
+            transition: $transition;
             display: flex;
-            flex-direction: column;
-
-            .message {
-                padding: 10px 15px;
-                margin-bottom: 10px;
-                border-radius: $border-radius;
-                background-color: lighten($background-color, 5%);
-                box-shadow: 0 2px 5px $shadow-color;
-                align-items: flex-start;
-                max-width: 90%;
-                width: fit-content;
-                word-break: break-word;
-                position: relative;
-
-                .citation-button {
-                    background-color: $citation-color;
-                    color: #fff;
-                    border: none;
-                    border-radius: 50%;
-                    cursor: pointer;
-                    width: 20px;
-                    height: 20px;
-                    display: inline-flex;
-                    align-items: center;
-                    justify-content: center;
-                    font-size: 12px;
-                    font-weight: bold;
-                    margin: 0 2px;
-                    padding: 0;
-                    transition: background-color 0.3s;
-
-                    &:hover {
-                        background-color: $citation-hover-color;
-                    }
-                }
-
-                &.user {
-                    align-self: flex-end;
-                    background-color: $button-color;
-                    color: #fff;
-                }
-
-                &.chatbot {
-                    align-self: flex-start;
-                    background-color: $input-background;
-                    color: $text-color;
-                }
-
-                span {
-                    flex-grow: 1;
-                    padding-right: 10px;
-                }
-
-                img {
-                    max-width: 50px;
-                    max-height: 50px;
-                    border-radius: 50%;
-                }
+            align-items: center;
+            justify-content: center;
+
+            &:hover {
+                background-color: darken($primary-color, 10%);
+            }
+
+            &:disabled {
+                background-color: $light-text-color;
+                cursor: not-allowed;
             }
 
-            .follow-up-questions {
-                margin-top: 10px;
-                width: 100%;
-
-                h4 {
-                    margin-bottom: 5px;
-                    font-size: 14px;
-                }
-
-                .follow-up-button {
-                    background-color: $follow-up-bg-color;
-                    border: 1px solid #ddd;
-                    border-radius: 8px;
-                    padding: 8px 10px;
-                    margin: 4px 0;
-                    cursor: pointer;
-                    transition: background-color 0.3s;
-                    display: block;
-                    width: 100%;
-                    text-align: left;
-                    white-space: normal;
-                    word-wrap: break-word;
-                    font-size: 12px;
-                    color: $text-color;
-                    min-height: 40px;
-                    height: auto;
-                    line-height: 1.3;
-
-                    &:hover {
-                        background-color: $follow-up-hover-bg-color;
-                    }
-                }
+            .spinner {
+                height: 24px;
+                width: 24px;
+                border: 3px solid rgba(255, 255, 255, 0.3);
+                border-top: 3px solid #fff;
+                border-radius: 50%;
+                animation: spin 1s linear infinite;
             }
         }
     }
+}
 
-    .chat-form {
-        display: flex;
-        flex-grow: 0;
-        bottom: 0;
-        width: 100%;
-        padding: 10px;
-        background-color: $input-background;
-        box-shadow: inset 0 -1px 2px $shadow-color;
-        margin-bottom: 0;
+.message {
+    max-width: 80%;
+    margin-bottom: 20px;
+    padding: 16px 20px;
+    border-radius: 18px;
+    font-size: 15px;
+    line-height: 1.5;
+    box-shadow: 0 2px 4px $shadow-color;
+
+    &.user {
+        align-self: flex-end;
+        background-color: $primary-color;
+        color: white;
+        border-bottom-right-radius: 4px;
+    }
 
-        input[type='text'] {
-            flex-grow: 1;
-            border: 1px solid darken($input-background, 10%);
-            border-radius: $border-radius;
-            padding: 8px 12px;
-            margin-right: 10px;
+    &.chatbot {
+        align-self: flex-start;
+        background-color: $secondary-color;
+        color: $text-color;
+        border-bottom-left-radius: 4px;
+    }
+
+    .toggle-info {
+        background-color: transparent;
+        color: $primary-color;
+        border: 1px solid $primary-color;
+        width: 100%;
+        height: fit-content;
+        border-radius: 8px;
+        padding: 10px 16px;
+        font-size: 14px;
+        cursor: pointer;
+        transition: $transition;
+        margin-top: 10px;
+
+        &:hover {
+            background-color: rgba($primary-color, 0.1);
         }
+    }
+}
 
-        button {
-            padding: 8px 16px;
-            background-color: $button-color;
-            color: #fff;
-            border: none;
-            border-radius: $border-radius;
-            cursor: pointer;
-            transition: background-color 0.3s;
-            min-width: 80px;
+.follow-up-questions {
+    margin-top: 15px;
 
-            &:hover {
-                background-color: $button-hover-color;
-            }
+    h4 {
+        font-size: 15px;
+        font-weight: 600;
+        margin-bottom: 10px;
+    }
+
+    .questions-list {
+        display: flex;
+        flex-direction: column;
+        gap: 10px;
+    }
+
+    .follow-up-button {
+        background-color: #fff;
+        color: $primary-color;
+        border: 1px solid $primary-color;
+        border-radius: 8px;
+        padding: 10px 16px;
+        font-size: 14px;
+        cursor: pointer;
+        transition: $transition;
+        text-align: left;
+        white-space: normal;
+        word-wrap: break-word;
+        width: 100%;
+        height: fit-content;
+
+        &:hover {
+            background-color: $primary-color;
+            color: #fff;
         }
     }
 }
 
-.uploading-overlay {
-    position: absolute;
-    top: 0;
-    left: 0;
-    width: 100%;
-    height: 100%;
-    background-color: rgba($background-color, 0.95);
-    display: flex;
-    justify-content: center;
+.citation-button {
+    display: inline-flex;
     align-items: center;
-    font-size: 1.5em;
+    justify-content: center;
+    width: 20px;
+    height: 20px;
+    border-radius: 50%;
+    background-color: rgba(0, 0, 0, 0.1);
     color: $text-color;
-    z-index: 10;
-
-    &::before {
-        content: 'Uploading Docs...';
-        font-weight: bold;
+    font-size: 12px;
+    font-weight: bold;
+    margin-left: 5px;
+    cursor: pointer;
+    transition: $transition;
+    vertical-align: middle;
+
+    &:hover {
+        background-color: rgba(0, 0, 0, 0.2);
     }
 }
 
-.modal {
-    position: fixed;
+.uploading-overlay {
+    position: absolute;
     top: 0;
     left: 0;
-    width: 100%;
-    height: 100%;
+    right: 0;
+    bottom: 0;
+    background-color: rgba(255, 255, 255, 0.8);
     display: flex;
     justify-content: center;
     align-items: center;
-    background-color: rgba(0, 0, 0, 0.4);
-
-    .modal-content {
-        background-color: $input-background;
-        color: $text-color;
-        padding: 20px;
-        border-radius: $border-radius;
-        box-shadow: 0 2px 10px $shadow-color;
-        display: flex;
-        flex-direction: column;
-        align-items: center;
-        width: auto;
-        min-width: 300px;
-
-        h4 {
-            margin-bottom: 15px;
-        }
-
-        p {
-            margin-bottom: 20px;
-        }
+    z-index: 1000;
+}
 
-        button {
-            padding: 10px 20px;
-            background-color: $button-color;
-            color: #fff;
-            border: none;
-            border-radius: $border-radius;
-            cursor: pointer;
-            transition: background-color 0.3s;
+@keyframes spin {
+    0% {
+        transform: rotate(0deg);
+    }
+    100% {
+        transform: rotate(360deg);
+    }
+}
 
-            &:hover {
-                background-color: $button-hover-color;
-            }
-        }
+@media (max-width: 768px) {
+    .chat-box {
+        border-radius: 0;
     }
-    .thought-text {
-        color: #6c757d;
-        font-style: italic;
+
+    .message {
+        max-width: 90%;
     }
 }
diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index 36416a330..4a98f8dc1 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -21,6 +21,8 @@ import { DocumentManager } from '../../../util/DocumentManager';
 import { v4 as uuidv4 } from 'uuid';
 import { chunk } from 'lodash';
 import { DocUtils } from '../../../documents/DocUtils';
+import { createRef } from 'react';
+import { ClientUtils } from '../../../../ClientUtils';
 
 dotenv.config();
 
@@ -37,10 +39,10 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     @observable private linked_csv_files: { filename: string; id: string; text: string }[] = [];
     private openai: OpenAI;
     private vectorstore_id: string;
-    private documents: AI_Document[] = [];
-    private _oldWheel: any;
     private vectorstore: Vectorstore;
     private agent: Agent; // Add the ChatBot instance
+    private _oldWheel: HTMLDivElement | null = null;
+    private messagesRef: React.RefObject<HTMLDivElement>;
 
     public static LayoutString(fieldKey: string) {
         return FieldView.LayoutString(ChatBox, fieldKey);
@@ -59,6 +61,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         }
         this.vectorstore = new Vectorstore(this.vectorstore_id, this.retrieveDocIds);
         this.agent = new Agent(this.vectorstore, this.retrieveSummaries, this.retrieveFormattedHistory, this.retrieveCSVData, this.addLinkedUrlDoc);
+        this.messagesRef = React.createRef<HTMLDivElement>();
 
         reaction(
             () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, content: msg.content, follow_up_questions: msg.follow_up_questions, citations: msg.citations })),
@@ -133,6 +136,23 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         return new OpenAI(configuration);
     }
 
+    addScrollListener = () => {
+        if (this.messagesRef.current) {
+            this.messagesRef.current.addEventListener('wheel', this.onPassiveWheel, { passive: false });
+        }
+    };
+
+    removeScrollListener = () => {
+        if (this.messagesRef.current) {
+            this.messagesRef.current.removeEventListener('wheel', this.onPassiveWheel);
+        }
+    };
+
+    scrollToBottom = () => {
+        if (this.messagesRef.current) {
+        }
+    };
+
     onPassiveWheel = (e: WheelEvent) => {
         if (this._props.isContentActive()) {
             e.stopPropagation();
@@ -160,6 +180,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                             this.current_message = { ...this.current_message, processing_info: update };
                         }
                     });
+                    this.scrollToBottom();
                 };
 
                 const finalMessage = await this.agent.askAgent(trimmedText, onUpdate);
@@ -176,8 +197,10 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 this.history.push({ role: ASSISTANT_ROLE.ASSISTANT, content: [{ index: 0, type: TEXT_TYPE.ERROR, text: 'Sorry, I encountered an error while processing your request.', citation_ids: null }], processing_info: [] });
             } finally {
                 this.isLoading = false;
+                this.scrollToBottom();
             }
         }
+        this.scrollToBottom();
     };
 
     @action
@@ -202,6 +225,11 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         doc.chunk_simpl = JSON.stringify({ chunks: [chunkToAdd] });
     };
 
+    @computed
+    get userName() {
+        return ClientUtils.CurrentUserEmail;
+    }
+
     @action
     handleCitationClick = (citation: Citation) => {
         console.log('Citation clicked:', citation);
@@ -276,6 +304,10 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         return highlight_doc;
     };
 
+    componentDidUpdate() {
+        this.scrollToBottom();
+    }
+
     componentDidMount() {
         this._props.setContentViewBox?.(this);
         if (this.dataDoc.data) {
@@ -332,6 +364,11 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 console.log('Deleted docs: ', change.oldValue);
             }
         });
+        this.addScrollListener();
+    }
+
+    componentWillUnmount() {
+        this.removeScrollListener();
     }
 
     @computed
@@ -411,35 +448,41 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     };
     render() {
         return (
-            <div className="chatBox">
-                {this.isUploadingDocs && <div className="uploading-overlay"></div>}
-                <div
-                    className="scroll-box chat-content"
-                    ref={r => {
-                        this._oldWheel?.removeEventListener('wheel', this.onPassiveWheel);
-                        this._oldWheel = r;
-                        r?.addEventListener('wheel', this.onPassiveWheel, { passive: false });
-                    }}>
-                    <div className="messages">
-                        {this.history.map((message, index) => (
-                            <MessageComponentBox key={index} message={message} index={index} onFollowUpClick={this.handleFollowUpClick} onCitationClick={this.handleCitationClick} updateMessageCitations={this.updateMessageCitations} />
-                        ))}
-                        {!this.current_message ? null : (
-                            <MessageComponentBox
-                                key={this.history.length}
-                                message={this.current_message}
-                                index={this.history.length}
-                                onFollowUpClick={this.handleFollowUpClick}
-                                onCitationClick={this.handleCitationClick}
-                                updateMessageCitations={this.updateMessageCitations}
-                            />
-                        )}
+            <div className="chat-box">
+                {this.isUploadingDocs && (
+                    <div className="uploading-overlay">
+                        <div className="spinner"></div>
                     </div>
+                )}
+                <div className="chat-header">
+                    <h2>{this.userName()}'s AI Assistant</h2>
                 </div>
-                <form onSubmit={this.askGPT} className="chat-form">
-                    <input type="text" name="messageInput" autoComplete="off" placeholder="Type a message..." value={this.inputValue} onChange={e => (this.inputValue = e.target.value)} />
-                    <button type="submit" disabled={this.isLoading}>
-                        {this.isLoading ? 'Thinking...' : 'Send'}
+                <div className="chat-messages" ref={this.messagesRef}>
+                    {this.history.map((message, index) => (
+                        <MessageComponentBox key={index} message={message} index={index} onFollowUpClick={this.handleFollowUpClick} onCitationClick={this.handleCitationClick} updateMessageCitations={this.updateMessageCitations} />
+                    ))}
+                    {this.current_message && (
+                        <MessageComponentBox
+                            key={this.history.length}
+                            message={this.current_message}
+                            index={this.history.length}
+                            onFollowUpClick={this.handleFollowUpClick}
+                            onCitationClick={this.handleCitationClick}
+                            updateMessageCitations={this.updateMessageCitations}
+                        />
+                    )}
+                </div>
+                <form onSubmit={this.askGPT} className="chat-input">
+                    <input type="text" name="messageInput" autoComplete="off" placeholder="Type your message here..." value={this.inputValue} onChange={e => (this.inputValue = e.target.value)} />
+                    <button className="submit-button" type="submit" disabled={this.isLoading}>
+                        {this.isLoading ? (
+                            <div className="spinner"></div>
+                        ) : (
+                            <svg viewBox="0 0 24 24" width="24" height="24" stroke="currentColor" strokeWidth="2" fill="none" strokeLinecap="round" strokeLinejoin="round">
+                                <line x1="22" y1="2" x2="11" y2="13"></line>
+                                <polygon points="22 2 15 22 11 13 2 9 22 2"></polygon>
+                            </svg>
+                        )}
                     </button>
                 </form>
             </div>
diff --git a/src/client/views/nodes/ChatBox/MessageComponent.scss b/src/client/views/nodes/ChatBox/MessageComponent.scss
deleted file mode 100644
index 6fcc0e5e7..000000000
--- a/src/client/views/nodes/ChatBox/MessageComponent.scss
+++ /dev/null
@@ -1,10 +0,0 @@
-MessageComponent-citation {
-    color: lightblue;
-    vertical-align: super;
-    font-size: smaller;
-}
-MessageComponent-file_path {
-    color: lightblue;
-    vertical-align: baseline;
-    font-size: inherit;
-}
diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index 0b8fa6b96..00e9795e3 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -1,7 +1,7 @@
-import React from 'react';
+import React, { useState } from 'react';
 import { observer } from 'mobx-react';
 import { AssistantMessage, Citation, MessageContent, PROCESSING_TYPE, ProcessingInfo, TEXT_TYPE } from './types';
-import Markdown from 'react-markdown';
+import ReactMarkdown from 'react-markdown';
 
 interface MessageComponentProps {
     message: AssistantMessage;
@@ -12,37 +12,20 @@ interface MessageComponentProps {
 }
 
 const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message, index, onFollowUpClick, onCitationClick, updateMessageCitations }) {
+    const [dropdownOpen, setDropdownOpen] = useState(false);
+
     const renderContent = (item: MessageContent) => {
         const i = item.index;
         if (item.type === TEXT_TYPE.GROUNDED) {
             const citation_ids = item.citation_ids || [];
             return (
                 <span key={i} className="grounded-text">
-                    {item.text}
+                    <ReactMarkdown>{item.text}</ReactMarkdown>
                     {citation_ids.map((id, idx) => {
                         const citation = message.citations?.find(c => c.citation_id === id);
                         if (!citation) return null;
                         return (
-                            <button
-                                key={i + idx}
-                                className="citation-button"
-                                onClick={() => onCitationClick(citation)}
-                                style={{
-                                    display: 'inline-flex',
-                                    alignItems: 'center',
-                                    justifyContent: 'center',
-                                    width: '20px',
-                                    height: '20px',
-                                    borderRadius: '50%',
-                                    border: 'none',
-                                    background: '#ff6347',
-                                    color: 'white',
-                                    fontSize: '12px',
-                                    fontWeight: 'bold',
-                                    cursor: 'pointer',
-                                    margin: '0 2px',
-                                    padding: 0,
-                                }}>
+                            <button key={i + idx} className="citation-button" onClick={() => onCitationClick(citation)}>
                                 {idx + 1}
                             </button>
                         );
@@ -52,49 +35,65 @@ const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message
         } else if (item.type === TEXT_TYPE.NORMAL) {
             return (
                 <span key={i} className="normal-text">
-                    {item.text}
+                    <ReactMarkdown>{item.text}</ReactMarkdown>
                 </span>
             );
         } else if ('query' in item) {
-            // Handle the case where the item has a query property
             return (
                 <span key={i} className="query-text">
-                    {JSON.stringify(item.query)}
+                    <ReactMarkdown>{JSON.stringify(item.query)}</ReactMarkdown>
                 </span>
             );
         } else {
-            // Fallback for any other unexpected cases
-            return <span key={i}>{JSON.stringify(item)}</span>;
+            return (
+                <span key={i}>
+                    <ReactMarkdown>{JSON.stringify(item)}</ReactMarkdown>
+                </span>
+            );
         }
     };
 
-    console.log(message.processing_info);
+    const hasProcessingInfo = message.processing_info && message.processing_info.length > 0;
+
+    const renderProcessingInfo = (info: ProcessingInfo) => {
+        if (info.type === PROCESSING_TYPE.THOUGHT) {
+            return (
+                <div key={info.index} className="dropdown-item">
+                    <strong>Thought:</strong> {info.content}
+                </div>
+            );
+        } else if (info.type === PROCESSING_TYPE.ACTION) {
+            return (
+                <div key={info.index} className="dropdown-item">
+                    <strong>Action:</strong> {info.content}
+                </div>
+            );
+        } else {
+            return null;
+        }
+    };
 
     return (
         <div className={`message ${message.role}`}>
-            {message.processing_info &&
-                (message.processing_info as ProcessingInfo[]).map(item =>
-                    item.type === PROCESSING_TYPE.THOUGHT ? (
-                        <div key={item.index} className="thought">
-                            <strong>Thought:</strong> {item.content}
-                        </div>
-                    ) : item.type === PROCESSING_TYPE.ACTION ? (
-                        <div key={item.index} className="action">
-                            <strong>Action:</strong> {item.content}
-                        </div>
-                    ) : (
-                        <div key={item.index} className="error"></div>
-                    )
-                )}
-            <div>{message.content && message.content.map(messageFragment => <React.Fragment key={messageFragment.index}>{renderContent(messageFragment)}</React.Fragment>)}</div>
+            <div className="message-content">{message.content && message.content.map(messageFragment => <React.Fragment key={messageFragment.index}>{renderContent(messageFragment)}</React.Fragment>)}</div>
+            {hasProcessingInfo && (
+                <div className="processing-info">
+                    <button className="toggle-info" onClick={() => setDropdownOpen(!dropdownOpen)}>
+                        {dropdownOpen ? 'Hide Agent Thoughts/Actions' : 'Show Agent Thoughts/Actions'}
+                    </button>
+                    {dropdownOpen && <div className="info-content">{message.processing_info.map(renderProcessingInfo)}</div>}
+                </div>
+            )}
             {message.follow_up_questions && message.follow_up_questions.length > 0 && (
                 <div className="follow-up-questions">
                     <h4>Follow-up Questions:</h4>
-                    {message.follow_up_questions.map((question, idx) => (
-                        <button key={idx} className="follow-up-button" onClick={() => onFollowUpClick(question)}>
-                            {question}
-                        </button>
-                    ))}
+                    <div className="questions-list">
+                        {message.follow_up_questions.map((question, idx) => (
+                            <button key={idx} className="follow-up-button" onClick={() => onFollowUpClick(question)}>
+                                {question}
+                            </button>
+                        ))}
+                    </div>
                 </div>
             )}
         </div>
diff --git a/src/server/ApiManagers/AssistantManager.ts b/src/server/ApiManagers/AssistantManager.ts
index cd26ca79b..9b85dbbe8 100644
--- a/src/server/ApiManagers/AssistantManager.ts
+++ b/src/server/ApiManagers/AssistantManager.ts
@@ -14,6 +14,8 @@ import { PartitionResponse } from 'unstructured-client/sdk/models/operations';
 import { ChunkingStrategy, Strategy } from 'unstructured-client/sdk/models/shared';
 import * as cheerio from 'cheerio';
 import { ScrapflyClient, ScrapeConfig } from 'scrapfly-sdk';
+import { google } from 'googleapis';
+import puppeteer from 'puppeteer';
 
 export enum Directory {
     parsed_files = 'parsed_files',
@@ -55,6 +57,7 @@ export default class AssistantManager extends ApiManager {
             },
         });
         const scrapflyClient = new ScrapflyClient({ key: process.env._CLIENT_SCRAPFLY_API_KEY! });
+        const customsearch = google.customsearch('v1');
 
         register({
             method: Method.POST,
@@ -89,20 +92,18 @@ export default class AssistantManager extends ApiManager {
             secureHandler: async ({ req, res }) => {
                 const { query } = req.body;
                 try {
-                    const response = await axios.get('http://api.serpstack.com/search', {
-                        params: {
-                            access_key: process.env._CLIENT_SERPSTACK_API_KEY,
-                            query: query,
-                        },
+                    const response = await customsearch.cse.list({
+                        q: query,
+                        cx: process.env._CLIENT_GOOGLE_SEARCH_ENGINE_ID,
+                        key: process.env._CLIENT_GOOGLE_API_KEY,
+                        safe: 'active',
                     });
-                    console.log(response.data);
 
-                    const results = response.data.organic_results.map((result: any) => ({
-                        url: result.url,
-                        snippet: result.snippet,
-                    }));
-
-                    console.log(results);
+                    const results =
+                        response.data.items?.map((item: any) => ({
+                            url: item.link,
+                            snippet: item.snippet,
+                        })) || [];
 
                     res.send({ results });
                 } catch (error: any) {
@@ -144,6 +145,7 @@ export default class AssistantManager extends ApiManager {
                     const scrapedImagesDirectory = pathToDirectory(Directory.scrape_images);
                     const filePath = serverPathToFile(Directory.scrape_images, url_filename);
 
+                    // Check if the image already exists
                     if (fs.existsSync(filePath)) {
                         const imageBuffer = await readFileAsync(filePath);
                         const base64Image = imageBuffer.toString('base64');
@@ -151,33 +153,27 @@ export default class AssistantManager extends ApiManager {
                         return res.send({ website_image_base64: base64Image });
                     }
 
+                    // Create the directory if it doesn't exist
                     if (!fs.existsSync(scrapedImagesDirectory)) {
                         fs.mkdirSync(scrapedImagesDirectory);
                     }
 
-                    const result = await scrapflyClient.scrape(
-                        new ScrapeConfig({
-                            url: url,
-                            render_js: true,
-                            screenshots: { everything: 'fullpage' },
-                        })
-                    );
-
-                    const screenshotPromises = Object.entries(result.result.screenshots).map(async ([name, screenshot]) => {
-                        const response = await axios.get(screenshot.url, {
-                            params: {
-                                key: process.env._CLIENT_SCRAPFLY_API_KEY!,
-                                options: 'print_media_format',
-                                proxy_pool: 'public_residential_pool',
-                            },
-                            responseType: 'arraybuffer',
-                        });
-                        await fs.promises.writeFile(filePath, response.data);
-                        return response.data.toString('base64');
+                    // Launch Puppeteer to take a screenshot of the webpage
+                    const browser = await puppeteer.launch({
+                        args: ['--no-sandbox', '--disable-setuid-sandbox'],
                     });
-
-                    const base64Screenshots = await Promise.all(screenshotPromises);
-                    res.send({ website_image_base64: base64Screenshots[0] });
+                    const page = await browser.newPage();
+                    await page.setUserAgent('Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36');
+                    await page.goto(url, { waitUntil: 'networkidle2' });
+                    const screenshotBuffer = await page.screenshot({ fullPage: true });
+                    await browser.close();
+
+                    // Save the screenshot to the file system
+                    await writeFileAsync(filePath, screenshotBuffer);
+
+                    // Return the base64-encoded image
+                    const base64Image = Buffer.from(screenshotBuffer).toString('base64');
+                    res.send({ website_image_base64: base64Image });
                 } catch (error: any) {
                     console.error('Error scraping website:', error);
                     res.status(500).send({ error: 'Failed to scrape website', details: error.message });