add Message.extra

Author: ngxson

examples/server/public/index.html.gz

@@ -159,6 +159,30 @@ export default function ChatMessage({
                         </div>
                       </details>
                     )}
+
+                    {msg.extra && msg.extra.length > 0 && (
+                      <details
+                        className={classNames({
+                          'collapse collapse-arrow mb-4 bg-base-200': true,
+                          'bg-opacity-10': msg.role !== 'assistant',
+                        })}
+                      >
+                        <summary className="collapse-title">
+                          Extra content
+                        </summary>
+                        <div className="collapse-content">
+                          {msg.extra.map((extra) =>
+                            extra.type === 'textFile' ? (
+                              <div key={extra.name}>
+                                <b>{extra.name}</b>
+                                <pre>{extra.content}</pre>
+                              </div>
+                            ) : null // TODO: support other extra types
+                          )}
+                        </div>
+                      </details>
+                    )}
+
                     <MarkdownDisplay
                       content={content}
                       isGenerating={isPending}

examples/server/webui/src/components/ChatMessage.tsx

@@ -5,6 +5,7 @@ import { CanvasType, Message, PendingMessage } from '../utils/types';
 import { classNames, throttle } from '../utils/misc';
 import CanvasPyInterpreter from './CanvasPyInterpreter';
 import StorageUtils from '../utils/storage';
+import { useVSCodeContext } from '../utils/llama-vscode';
 
 /**
  * A message display is a message node with additional information for rendering.
@@ -79,36 +80,16 @@ export default function ChatScreen() {
     pendingMessages,
     canvasData,
     replaceMessageAndGenerate,
-    setExtraContext
   } = useAppContext();
   const [inputMsg, setInputMsg] = useState('');
   const inputRef = useRef<HTMLTextAreaElement>(null);
 
-  // Accept setText message from a parent window and set inputMsg and extraContext
-  useEffect(() => {
-    const handleMessage = (event: MessageEvent) => {
-      if (event.data?.command === 'setText') {
-        setInputMsg(event.data?.text);
-        setExtraContext(event.data?.context)
-        inputRef.current?.focus();
-      }
-    };
-
-    window.addEventListener('message', handleMessage);
-    return () => window.removeEventListener('message', handleMessage);
-  }, []);
-
-  // Add a keydown listener that sends the "escapePressed" message to the parent window
-  useEffect(() => {
-    const handleKeyDown = (event: KeyboardEvent) => {
-      if (event.key === 'Escape') {
-          window.parent.postMessage({ command: 'escapePressed' }, '*');
-      }
-    };
-
-    window.addEventListener('keydown', handleKeyDown);
-    return () => window.removeEventListener('keydown', handleKeyDown);
-  }, []);
+  const { extraContext, clearExtraContext } = useVSCodeContext(
+    inputRef,
+    setInputMsg
+  );
+  // TODO: improve this when we have "upload file" feature
+  const currExtra: Message['extra'] = extraContext ? [extraContext] : undefined;
 
   // keep track of leaf node for rendering
   const [currNodeId, setCurrNodeId] = useState<number>(-1);
@@ -143,10 +124,20 @@ export default function ChatScreen() {
     setCurrNodeId(-1);
     // get the last message node
     const lastMsgNodeId = messages.at(-1)?.msg.id ?? null;
-    if (!(await sendMessage(currConvId, lastMsgNodeId, inputMsg, onChunk))) {
+    if (
+      !(await sendMessage(
+        currConvId,
+        lastMsgNodeId,
+        inputMsg,
+        currExtra,
+        onChunk
+      ))
+    ) {
       // restore the input message if failed
       setInputMsg(lastInpMsg);
     }
+    // OK
+    clearExtraContext();
   };
 
   const handleEditMessage = async (msg: Message, content: string) => {
@@ -157,6 +148,7 @@ export default function ChatScreen() {
       viewingChat.conv.id,
       msg.parent,
       content,
+      msg.extra,
       onChunk
     );
     setCurrNodeId(-1);
@@ -171,6 +163,7 @@ export default function ChatScreen() {
       viewingChat.conv.id,
       msg.parent,
       null,
+      msg.extra,
       onChunk
     );
     setCurrNodeId(-1);

examples/server/webui/src/components/ChatScreen.tsx

@@ -25,13 +25,15 @@ interface AppContextValue {
     convId: string | null,
     leafNodeId: Message['id'] | null,
     content: string,
+    extra: Message['extra'],
     onChunk: CallbackGeneratedChunk
   ) => Promise<boolean>;
   stopGenerating: (convId: string) => void;
   replaceMessageAndGenerate: (
     convId: string,
     parentNodeId: Message['id'], // the parent node of the message to be replaced
     content: string | null,
+    extra: Message['extra'],
     onChunk: CallbackGeneratedChunk
   ) => Promise<void>;
 
@@ -44,10 +46,6 @@ interface AppContextValue {
   saveConfig: (config: typeof CONFIG_DEFAULT) => void;
   showSettings: boolean;
   setShowSettings: (show: boolean) => void;
-
-  // extra context
-  extraContext: string;
-  setExtraContext: (extraCtx: string) => void;
 }
 
 // this callback is used for scrolling to the bottom of the chat and switching to the last node
@@ -86,7 +84,6 @@ export const AppContextProvider = ({
   const [config, setConfig] = useState(StorageUtils.getConfig());
   const [canvasData, setCanvasData] = useState<CanvasData | null>(null);
   const [showSettings, setShowSettings] = useState(false);
-  const [extraContext, setExtraContext] = useState("");
 
   // handle change when the convId from URL is changed
   useEffect(() => {
@@ -179,10 +176,6 @@ export const AppContextProvider = ({
           : [{ role: 'system', content: config.systemMessage } as APIMessage]),
         ...normalizeMsgsForAPI(currMessages),
       ];
-      if (extraContext && extraContext != ""){
-        // insert extra context just before the user messages
-        messages.splice(config.systemMessage.length === 0 ? 0 : 1, 0, { role: 'user', content:extraContext } as APIMessage)
-      }
       if (config.excludeThoughtOnReq) {
         messages = filterThoughtFromMsgs(messages);
       }
@@ -283,6 +276,7 @@ export const AppContextProvider = ({
     convId: string | null,
     leafNodeId: Message['id'] | null,
     content: string,
+    extra: Message['extra'],
     onChunk: CallbackGeneratedChunk
   ): Promise<boolean> => {
     if (isGenerating(convId ?? '') || content.trim().length === 0) return false;
@@ -307,6 +301,7 @@ export const AppContextProvider = ({
         convId,
         role: 'user',
         content,
+        extra,
         parent: leafNodeId,
         children: [],
       },
@@ -333,6 +328,7 @@ export const AppContextProvider = ({
     convId: string,
     parentNodeId: Message['id'], // the parent node of the message to be replaced
     content: string | null,
+    extra: Message['extra'],
     onChunk: CallbackGeneratedChunk
   ) => {
     if (isGenerating(convId)) return;
@@ -348,6 +344,7 @@ export const AppContextProvider = ({
           convId,
           role: 'user',
           content,
+          extra,
           parent: parentNodeId,
           children: [],
         },
@@ -380,8 +377,6 @@ export const AppContextProvider = ({
         saveConfig,
         showSettings,
         setShowSettings,
-        extraContext,
-        setExtraContext,
       }}
     >
       {children}

examples/server/webui/src/utils/app.context.tsx

@@ -0,0 +1,62 @@
+import { useEffect, useState } from 'react';
+import { MessageExtraTextFile } from './types';
+
+// Extra context when using llama.cpp WebUI from llama-vscode, inside an iframe
+// Ref: https://github.com/ggml-org/llama.cpp/pull/11940
+
+interface SetTextEvData {
+  text: string;
+  context: string;
+}
+
+/**
+ * To test it:
+ * window.postMessage({ command: 'setText', text: 'Spot the syntax error', context: 'def test()\n  return 123' }, '*');
+ */
+
+export const useVSCodeContext = (
+  inputRef: React.RefObject<HTMLTextAreaElement>,
+  setInputMsg: (text: string) => void
+) => {
+  const [extraContext, setExtraContext] = useState<MessageExtraTextFile | null>(
+    null
+  );
+  console.log({ extraContext });
+
+  // Accept setText message from a parent window and set inputMsg and extraContext
+  useEffect(() => {
+    const handleMessage = (event: MessageEvent) => {
+      if (event.data?.command === 'setText') {
+        const data: SetTextEvData = event.data;
+        setInputMsg(data?.text);
+        setExtraContext({
+          type: 'textFile',
+          name: '(context)', // TODO: set filename
+          content: data?.context ?? '(no context)',
+        });
+        inputRef.current?.focus();
+      }
+    };
+
+    window.addEventListener('message', handleMessage);
+    return () => window.removeEventListener('message', handleMessage);
+  }, []);
+
+  // Add a keydown listener that sends the "escapePressed" message to the parent window
+  useEffect(() => {
+    const handleKeyDown = (event: KeyboardEvent) => {
+      if (event.key === 'Escape') {
+        window.parent.postMessage({ command: 'escapePressed' }, '*');
+      }
+    };
+
+    window.addEventListener('keydown', handleKeyDown);
+    return () => window.removeEventListener('keydown', handleKeyDown);
+  }, []);
+
+  return {
+    extraContext,
+    // call once the user message is sent, to clear the extra context
+    clearExtraContext: () => setExtraContext(null),
+  };
+};

examples/server/webui/src/utils/llama-vscode.ts

@@ -53,12 +53,24 @@ export const copyStr = (textToCopy: string) => {
 
 /**
  * filter out redundant fields upon sending to API
+ * also format extra into text
  */
 export function normalizeMsgsForAPI(messages: Readonly<Message[]>) {
   return messages.map((msg) => {
+    let newContent = '';
+
+    for (const extra of msg.extra ?? []) {
+      if (extra.type === 'textFile') {
+        // TODO: allow user to customize this via Settings
+        newContent += `\n===\nExtra file: ${extra.name}\n${extra.content}\n\n===\n`;
+      }
+    }
+
+    newContent += msg.content;
+
     return {
       role: msg.role,
-      content: msg.content,
+      content: newContent,
     };
   }) as APIMessage[];
 }

examples/server/webui/src/utils/misc.ts

@@ -42,11 +42,20 @@ export interface Message {
   role: 'user' | 'assistant' | 'system';
   content: string;
   timings?: TimingReport;
+  extra?: MessageExtra[];
   // node based system for branching
   parent: Message['id'];
   children: Message['id'][];
 }
 
+type MessageExtra = MessageExtraTextFile; // TODO: will add more in the future
+
+export interface MessageExtraTextFile {
+  type: 'textFile';
+  name: string;
+  content: string;
+}
+
 export type APIMessage = Pick<Message, 'role' | 'content'>;
 
 export interface Conversation {