fixing to work with python API

also added follow up questions

1 files changed, 41 insertions, 12 deletions

diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index ef6ce83b5..15c0811fb 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -14,24 +14,29 @@ interface MessageComponentProps {
     showModal: () => void;
     goToLinkedDoc: (url: string) => void;
     setCurrentFile: (file: { url: string }) => void;
+    onFollowUpClick: (question: string) => void; // New prop
     isCurrent?: boolean;
 }
 
-const MessageComponent: React.FC<MessageComponentProps> = function ({ message, toggleToolLogs, expandedLogIndex, goToLinkedDoc, index, showModal, setCurrentFile, isCurrent = false }) {
-    // const messageClass = `${message.role} ${isCurrent ? 'current-message' : ''}`;
-
+const MessageComponent: React.FC<MessageComponentProps> = function ({
+    message,
+    toggleToolLogs,
+    expandedLogIndex,
+    goToLinkedDoc,
+    index,
+    showModal,
+    setCurrentFile,
+    onFollowUpClick, // New prop
+    isCurrent = false,
+}) {
     const LinkRenderer = ({ href, children }: { href: string; children: React.ReactNode }) => {
-        // console.log(href + " " + children)
         const regex = /([a-zA-Z0-9_.!-]+)~~~(citation|file_path)/;
         const matches = href.match(regex);
-        // console.log(href)
-        // console.log(matches)
         const url = matches ? matches[1] : href;
         const linkType = matches ? matches[2] : null;
         if (linkType === 'citation') {
             children = <TbInfoCircleFilled />;
         }
-        // console.log(linkType)
         const style = {
             color: 'lightblue',
             verticalAlign: linkType === 'citation' ? 'super' : 'baseline',
@@ -56,14 +61,38 @@ const MessageComponent: React.FC<MessageComponentProps> = function ({ message, t
         );
     };
 
+    const parseMessage = (text: string) => {
+        const answerMatch = text.match(/<answer>([\s\S]*?)<\/answer>/);
+        const followUpMatch = text.match(/<follow_up_question>([\s\S]*?)<\/follow_up_question>/);
+
+        const answer = answerMatch ? answerMatch[1] : text;
+        const followUpQuestions = followUpMatch
+            ? followUpMatch[1]
+                  .split('\n')
+                  .filter(q => q.trim())
+                  .map(q => q.replace(/^\d+\.\s*/, '').trim())
+            : [];
+
+        return { answer, followUpQuestions };
+    };
+
+    const { answer, followUpQuestions } = parseMessage(message.text);
+    console.log('Parsed answer:', answer);
+    console.log('Parsed follow-up questions:', followUpQuestions);
     return (
         <div className={`message ${message.role}`}>
-            <MathJaxContext>
-                <MathJax dynamic hideUntilTypeset="every">
-                    <ReactMarkdown components={{ a: LinkRenderer }}>{message.text ? message.text : ''}</ReactMarkdown>
-                </MathJax>
-            </MathJaxContext>
+            <ReactMarkdown components={{ a: LinkRenderer }}>{answer}</ReactMarkdown>
             {message.image && <img src={message.image} alt="" />}
+            {followUpQuestions.length > 0 && (
+                <div className="follow-up-questions">
+                    <h4>Follow-up Questions:</h4>
+                    {followUpQuestions.map((question, idx) => (
+                        <button key={idx} className="follow-up-button" onClick={() => onFollowUpClick(question)}>
+                            {question}
+                        </button>
+                    ))}
+                </div>
+            )}
             <div className="message-footer">
                 {message.tool_logs && (
                     <button className="toggle-logs-button" onClick={() => toggleToolLogs(index)}>