Port aap_chatbot changes to ansible_ai_connect_chatbot

TamiTakamiya · TamiTakamiya · commit 161dc1acdc0b · 2025-04-03T16:23:20.000-04:00
diff --git a/ansible_ai_connect_chatbot/src/AnsibleChatbot/AnsibleChatbot.tsx b/ansible_ai_connect_chatbot/src/AnsibleChatbot/AnsibleChatbot.tsx
@@ -6,6 +6,7 @@ import {
   DropdownList,
   DropdownItem,
   DropdownGroup,
+  ExpandableSection,
 } from "@patternfly/react-core";
 
 import ChatbotContent from "@patternfly/chatbot/dist/dynamic/ChatbotContent";
@@ -44,6 +45,7 @@ import {
   ChatbotAlert,
   ChatbotConversationHistoryNav,
   ChatbotDisplayMode,
+  ChatbotFootnoteProps,
   ChatbotHeaderMain,
   ChatbotHeaderMenu,
   ChatbotHeaderSelectorDropdown,
@@ -58,7 +60,7 @@ import {
 } from "../Constants";
 import { SystemPromptModal } from "../SystemPromptModal/SystemPromptModal";
 
-const footnoteProps = {
+const footnoteProps: ChatbotFootnoteProps = {
   label: FOOTNOTE_LABEL,
   popover: {
     title: FOOTNOTE_TITLE,
@@ -346,6 +348,7 @@ export const AnsibleChatbot: React.FunctionComponent = () => {
                       {
                         referenced_documents,
                         scrollToHere,
+                        collapse,
                         ...message
                       }: ExtendedMessage,
                       index,
@@ -357,14 +360,27 @@ export const AnsibleChatbot: React.FunctionComponent = () => {
                             ref={messagesEndRef}
                           />
                         )}
-                        <div key={`m_div_${index}`}>
-                          <Message key={`m_msg_${index}`} {...message} />
-                          <ReferencedDocuments
-                            key={`m_docs_${index}`}
-                            caption="Refer to the following for more information:"
-                            referenced_documents={referenced_documents}
-                          />
-                        </div>
+                        {collapse ? (
+                          <div key={`m_div_${index}`}>
+                            <ExpandableSection toggleText="Show more">
+                              <Message key={`m_msg_${index}`} {...message} />
+                              <ReferencedDocuments
+                                key={`m_docs_${index}`}
+                                caption="Refer to the following for more information:"
+                                referenced_documents={referenced_documents}
+                              />
+                            </ExpandableSection>
+                          </div>
+                        ) : (
+                          <div key={`m_div_${index}`}>
+                            <Message key={`m_msg_${index}`} {...message} />
+                            <ReferencedDocuments
+                              key={`m_docs_${index}`}
+                              caption="Refer to the following for more information:"
+                              referenced_documents={referenced_documents}
+                            />
+                          </div>
+                        )}
                       </div>
                     ),
                   )}
diff --git a/ansible_ai_connect_chatbot/src/App.test.tsx b/ansible_ai_connect_chatbot/src/App.test.tsx
@@ -197,21 +197,67 @@ function mockFetchEventSource() {
     },
   ];
 
+  const streamAgentNormalData: object[] = [
+    {
+      event: "start",
+      data: { conversation_id: "6e33a46e-2b15-4128-8e5c-c6f637ebfbb4" },
+    },
+    {
+      event: "token",
+      data: { id: 0, token: "Let me search for " },
+    },
+    {
+      event: "token",
+      data: { id: 1, token: "information about EDA." },
+    },
+    {
+      event: "tool_call",
+      data: { id: 2, token: '{ "key":"value"}' },
+    },
+    {
+      event: "step_details",
+      data: { id: 3, token: '{ "key":"value"}' },
+    },
+    {
+      event: "token",
+      data: { id: 4, token: "EDA stands for " },
+    },
+    {
+      event: "token",
+      data: { id: 5, token: "Event Driven Ansible." },
+    },
+    {
+      event: "step_complete",
+      data: { id: 6, token: '{ "key":"value"}' },
+    },
+    {
+      event: "turn_complete",
+      data: { id: 0, token: "Turn complete." },
+    },
+  ];
+
   return vi.fn(async (_, init) => {
     let status = 200;
     let errorCase = false;
+    let agent = false;
     const o = JSON.parse(init.body);
     if (o.query.startsWith("status=")) {
       status = parseInt(o.query.substring(7));
     } else if (o.query.startsWith("error in stream")) {
       errorCase = true;
+    } else if (o.query.startsWith("agent")) {
+      agent = true;
     }
     console.log(`status ${status}`);
 
     const ok = status === 200;
     await init.onopen({ status, ok });
     if (status === 200) {
-      const streamData = errorCase ? streamErrorData : streamNormalData;
+      const streamData = agent
+        ? streamAgentNormalData
+        : errorCase
+          ? streamErrorData
+          : streamNormalData;
       for (const data of streamData) {
         init.onmessage({ data: JSON.stringify(data) });
       }
@@ -293,15 +339,15 @@ test("Basic chatbot interaction", async () => {
     .not.toBeInTheDocument();
 
   const footNoteLink = page.getByText(
-    "Lightspeed uses AI. Check for mistakes.",
+    "Always review AI-generated content prior to use.",
   );
   await footNoteLink.click();
   await expect
     .element(page.getByText("While Lightspeed strives for accuracy,"))
     .toBeVisible();
   await page.getByText("Got it").click();
   await expect
-    .element(view.getByText("While Lightspeed strives for accuracy,"))
+    .element(page.getByText("While Lightspeed strives for accuracy,"))
     .not.toBeVisible();
 
   await textArea.fill("Tell me about Ansible.");
@@ -608,6 +654,12 @@ test("Test system prompt override", async () => {
 });
 
 test("Chat streaming test", async () => {
+  let ghIssueLinkSpy = 0;
+  vi.stubGlobal("open", () => {
+    ghIssueLinkSpy++;
+  });
+  mockAxios(200);
+
   const view = await renderApp(false, true);
   const textArea = page.getByLabelText("Send a message...");
   await textArea.fill("Hello");
@@ -621,6 +673,53 @@ test("Chat streaming test", async () => {
       ),
     )
     .toBeVisible();
+
+  const thumbsDownIcon = await screen.findByRole("button", {
+    name: "Bad response",
+  });
+  await thumbsDownIcon.click();
+
+  const sureButton = await screen.findByText("Sure!");
+  await expect.element(sureButton).toBeVisible();
+  await sureButton.click();
+
+  expect(ghIssueLinkSpy).toEqual(1);
+});
+
+test("Agent chat streaming test", async () => {
+  let ghIssueLinkSpy = 0;
+  vi.stubGlobal("open", () => {
+    ghIssueLinkSpy++;
+  });
+  mockAxios(200, false, false, referencedDocumentExample);
+
+  const view = await renderApp(false, true);
+  const textArea = page.getByLabelText("Send a message...");
+  await textArea.fill("agent test");
+
+  await userEvent.keyboard("{Enter}");
+
+  await expect.element(view.getByText("Turn complete")).toBeVisible();
+
+  const thumbsDownIcon = await screen.findByRole("button", {
+    name: "Bad response",
+  });
+  await thumbsDownIcon.click();
+
+  const sureButton = await screen.findByText("Sure!");
+  await expect.element(sureButton).toBeVisible();
+  await sureButton.click();
+
+  expect(ghIssueLinkSpy).toEqual(1);
+
+  await expect
+    .element(view.getByText("EDA stands for Event Driven Ansible."))
+    .not.toBeVisible();
+  const showMoreLink = await screen.findByRole("button", { name: "Show more" });
+  await showMoreLink.click();
+  await expect
+    .element(view.getByText("EDA stands for Event Driven Ansible."))
+    .toBeVisible();
 });
 
 test("Chat streaming error at API call", async () => {
diff --git a/ansible_ai_connect_chatbot/src/Constants.ts b/ansible_ai_connect_chatbot/src/Constants.ts
@@ -13,7 +13,8 @@ export const TOO_MANY_REQUESTS_MSG =
   "_Chatbot service is busy with too many requests. Please try again later._";
 
 /* Footnote label */
-export const FOOTNOTE_LABEL = "Lightspeed uses AI. Check for mistakes.";
+export const FOOTNOTE_LABEL =
+  "Always review AI-generated content prior to use.";
 
 /* Footnote title */
 export const FOOTNOTE_TITLE = "Verify accuracy";
diff --git a/ansible_ai_connect_chatbot/src/types/Message.ts b/ansible_ai_connect_chatbot/src/types/Message.ts
@@ -37,6 +37,7 @@ export type ReferencedDocumentsProp = {
 export type ExtendedMessage = MessageProps & {
   referenced_documents: ReferencedDocument[];
   scrollToHere?: boolean;
+  collapse?: boolean;
 };
 
 export type ChatFeedback = {
diff --git a/ansible_ai_connect_chatbot/src/useChatbot/useChatbot.ts b/ansible_ai_connect_chatbot/src/useChatbot/useChatbot.ts
@@ -172,11 +172,11 @@ export const useChatbot = () => {
     });
   };
 
-  const appendMessageChunk = (chunk: string) => {
+  const appendMessageChunk = (chunk: string, query: string = "") => {
     setMessages((msgs: ExtendedMessage[]) => {
       const lastMessage = msgs[msgs.length - 1];
       if (!lastMessage || lastMessage.role === "user") {
-        const newMessage: ExtendedMessage = botMessage(chunk);
+        const newMessage: ExtendedMessage = botMessage(chunk, query);
         chunk = "";
         return [...msgs, newMessage];
       } else {
@@ -222,16 +222,34 @@ export const useChatbot = () => {
       timestamp: getTimestamp(),
       referenced_documents: [],
     };
-    const sendFeedback = async (sentiment: Sentiment) => {
-      if (typeof response === "object") {
+    const sendFeedback = async (
+      sentiment: Sentiment,
+      content: string = "",
+      referenced_documents: ReferencedDocument[] = [],
+    ) => {
+      if (typeof response === "string") {
+        const resp = {
+          conversation_id: conversationId
+            ? conversationId
+            : "00000000-0000-0000-0000-000000000000",
+          response: content,
+          referenced_documents,
+          truncated: false,
+        };
+        handleFeedback({ query, response: resp, sentiment, message });
+      } else {
         handleFeedback({ query, response, sentiment, message });
       }
     };
 
     message.actions = {
       positive: {
         onClick: () => {
-          sendFeedback(Sentiment.THUMBS_UP);
+          sendFeedback(
+            Sentiment.THUMBS_UP,
+            message.content,
+            message.referenced_documents,
+          );
           if (message.actions) {
             message.actions.positive.isDisabled = true;
             message.actions.negative.isDisabled = true;
@@ -241,7 +259,11 @@ export const useChatbot = () => {
       },
       negative: {
         onClick: () => {
-          sendFeedback(Sentiment.THUMBS_DOWN);
+          sendFeedback(
+            Sentiment.THUMBS_DOWN,
+            message.content,
+            message.referenced_documents,
+          );
           if (message.actions) {
             message.actions.positive.isDisabled = true;
             message.actions.negative.isDisabled = true;
@@ -323,10 +345,10 @@ export const useChatbot = () => {
     setAbortController(new AbortController());
   };
 
-  const handleSend = async (message: string | number) => {
+  const handleSend = async (query: string | number) => {
     const userMessage: ExtendedMessage = {
       role: "user",
-      content: message.toString(),
+      content: query.toString(),
       name: userName,
       avatar: userLogo,
       timestamp: getTimestamp(),
@@ -336,7 +358,7 @@ export const useChatbot = () => {
 
     const chatRequest: ChatRequest = {
       conversation_id: conversationId,
-      query: message.toString(),
+      query: query.toString(),
     };
 
     if (systemPrompt !== QUERY_SYSTEM_INSTRUCTION) {
@@ -385,8 +407,13 @@ export const useChatbot = () => {
                 });
               }
             },
-            onmessage(event: any) {
-              const message = JSON.parse(event.data);
+            onmessage(msg: any) {
+              let message = msg;
+              if (!msg.event) {
+                message = JSON.parse(msg.data);
+              } else {
+                message.data = JSON.parse(msg.data);
+              }
               if (message.event === "start") {
                 if (!conversationId) {
                   setConversationId(message.data.conversation_id);
@@ -395,7 +422,7 @@ export const useChatbot = () => {
                 if (message.data.token !== "") {
                   setIsLoading(false);
                 }
-                appendMessageChunk(message.data.token);
+                appendMessageChunk(message.data.token, query.toString());
               } else if (message.event === "end") {
                 if (message.data.referenced_documents.length > 0) {
                   addReferencedDocuments(message.data.referenced_documents);
@@ -409,6 +436,27 @@ export const useChatbot = () => {
                     `cause="${data.cause}"`,
                   variant: "danger",
                 });
+              } else if (
+                message.event === "tool_call" ||
+                message.event === "step_complete"
+              ) {
+                console.log(
+                  `!![${message.event}] ${JSON.stringify(message.data)}`,
+                );
+                appendMessageChunk(
+                  "\n\n`[" +
+                    message.event +
+                    "]`\n```json\n" +
+                    message.data.token +
+                    "\n```\n",
+                );
+              } else if (message.event === "turn_complete") {
+                setMessages((msgs: ExtendedMessage[]) => {
+                  const lastMessage = msgs[msgs.length - 1];
+                  lastMessage.collapse = true;
+                  msgs.push(botMessage(message.data.token, query.toString()));
+                  return msgs;
+                });
               }
             },
             onclose() {
@@ -440,10 +488,7 @@ export const useChatbot = () => {
           if (!conversationId) {
             setConversationId(chatResponse.conversation_id);
           }
-          const newBotMessage: any = botMessage(
-            chatResponse,
-            message.toString(),
-          );
+          const newBotMessage: any = botMessage(chatResponse, query.toString());
           newBotMessage.referenced_documents = referenced_documents;
           addMessage(newBotMessage);
         } else {
