From 9e7fb8377de1f3f1c29cbedd20a39772e3eafe89 Mon Sep 17 00:00:00 2001
From: Tiberiu Ichim <tiberiu.ichim@gmail.com>
Date: Tue, 17 Dec 2024 16:09:07 +0200
Subject: [PATCH] Show tool calls

---
 src/ChatBlock/ChatMessageBubble.jsx | 33 ++++++++++++++++++++++++-----
 src/ChatBlock/ChatWindow.jsx        | 20 +++++++++++++++--
 src/ChatBlock/schema.js             |  6 ++++++
 src/ChatBlock/style.less            |  5 +++++
 src/ChatBlock/useBackendChat.js     |  2 ++
 5 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/src/ChatBlock/ChatMessageBubble.jsx b/src/ChatBlock/ChatMessageBubble.jsx
index 11a7c72..bfc6bac 100644
--- a/src/ChatBlock/ChatMessageBubble.jsx
+++ b/src/ChatBlock/ChatMessageBubble.jsx
@@ -1,12 +1,15 @@
-import React from 'react';
 import loadable from '@loadable/component';
+import React from 'react';
+
 import { Citation } from './Citation';
 import { SourceDetails } from './Source';
-import { transformEmailsToLinks, SVGIcon } from './utils';
+import { SVGIcon, transformEmailsToLinks } from './utils';
 
 import BotIcon from './../icons/bot.svg';
 import UserIcon from './../icons/user.svg';
 
+const CITATION_MATCH = /\[\d+\](?![[(\])])/gm;
+
 const Markdown = loadable(() => import('react-markdown'));
 
 const components = (message) => {
@@ -45,8 +48,6 @@ const components = (message) => {
   };
 };
 
-const CITATION_MATCH = /\[\d+\](?![[(\])])/gm;
-
 function addCitations(text) {
   return text.replaceAll(CITATION_MATCH, (match) => {
     const number = match.match(/\d+/)[0];
@@ -54,8 +55,26 @@ function addCitations(text) {
   });
 }
 
+function ToolCall({ tool_args, tool_name, tool_result }) {
+  if (tool_name === 'run_search') {
+    return (
+      <div className="tool_info">
+        Searched for: <em>{tool_args?.query || ''}</em>
+      </div>
+    );
+  }
+  return null;
+}
+
 export function ChatMessageBubble(props) {
-  const { message, isLoading, isMostRecent, libs, onChoice } = props;
+  const {
+    message,
+    isLoading,
+    isMostRecent,
+    libs,
+    onChoice,
+    showToolCalls,
+  } = props;
   const { remarkGfm } = libs; // , rehypePrism
   const { citations = {}, documents, type } = message;
   const isUser = type === 'user';
@@ -95,6 +114,10 @@ export function ChatMessageBubble(props) {
         {icon}
 
         <div>
+          {showToolCalls &&
+            message.toolCalls?.map((info, index) => (
+              <ToolCall key={index} {...info} />
+            ))}
           <Markdown
             components={components(message)}
             remarkPlugins={[remarkGfm]}
diff --git a/src/ChatBlock/ChatWindow.jsx b/src/ChatBlock/ChatWindow.jsx
index 6aadd5e..6d753bf 100644
--- a/src/ChatBlock/ChatWindow.jsx
+++ b/src/ChatBlock/ChatWindow.jsx
@@ -21,7 +21,13 @@ function ChatWindow({
   isEditMode,
   ...data
 }) {
-  const { height, qgenAsistantId, enableQgen, scrollToInput } = data;
+  const {
+    height,
+    qgenAsistantId,
+    enableQgen,
+    scrollToInput,
+    showToolCalls,
+  } = data;
   const libs = { rehypePrism, remarkGfm }; // rehypePrism, remarkGfm
   const { onSubmit, messages, isStreaming, clearChat } = useBackendChat({
     persona,
@@ -62,7 +68,7 @@ function ChatWindow({
   }, [messages]);
 
   //eslint-disable-next-line
-  console.log(messages);
+  // console.log(messages);
 
   useScrollonStream({
     isStreaming,
@@ -73,6 +79,15 @@ function ChatWindow({
     debounce: 100, // time for debouncing
   });
 
+  // const timeoutRef = React.useRef();
+
+  // const debouncedSetInput = React.useCallback((e) => {
+  //   if (timeoutRef.current) clearTimeout(timeoutRef.current);
+  //   timeoutRef.current = setTimeout(() => {
+  //     setInput(e.target.value);
+  //   }, 1000);
+  // }, []);
+
   return (
     <div className="chat-window">
       <div className="messages">
@@ -107,6 +122,7 @@ function ChatWindow({
                   onChoice={(message) => {
                     onSubmit({ message });
                   }}
+                  showToolCalls={showToolCalls}
                 />
               ))}
               <div ref={endDivRef} /> {/* End div to mark the bottom */}
diff --git a/src/ChatBlock/schema.js b/src/ChatBlock/schema.js
index 96b01a3..81d6a76 100644
--- a/src/ChatBlock/schema.js
+++ b/src/ChatBlock/schema.js
@@ -17,6 +17,7 @@ export function ChatBlockSchema({ assistants }) {
           'height',
           'enableQgen',
           'scrollToInput',
+          'showToolCalls',
           'showAssistantTitle',
           'showAssistantDescription',
           'showAssistantPrompts',
@@ -39,6 +40,11 @@ export function ChatBlockSchema({ assistants }) {
         type: 'boolean',
         default: false,
       },
+      showToolCalls: {
+        title: 'Show query used in retriever',
+        type: 'boolean',
+        default: true,
+      },
       placeholderPrompt: {
         default: 'Ask a question',
         title: 'Prompt',
diff --git a/src/ChatBlock/style.less b/src/ChatBlock/style.less
index 55be362..9d74e4b 100644
--- a/src/ChatBlock/style.less
+++ b/src/ChatBlock/style.less
@@ -52,6 +52,11 @@
     i.icons {
       font-size: 1em;
     }
+
+    .tool_info {
+      font-size: small;
+      padding-bottom: 0.2em;
+    }
   }
 
   .sources {
diff --git a/src/ChatBlock/useBackendChat.js b/src/ChatBlock/useBackendChat.js
index ca98f4f..ad47855 100644
--- a/src/ChatBlock/useBackendChat.js
+++ b/src/ChatBlock/useBackendChat.js
@@ -447,6 +447,8 @@ export function useBackendChat({ persona, qgenAsistantId, enableQgen }) {
     setCurrChatSessionId(null);
   };
 
+  // console.log('history', messageHistory);
+
   return {
     messages: messageHistory,
     onSubmit: submitHandler.onSubmit,