[Query Copilot] Explanation bubble implementation (#1586)

* Explanation bubble implementation

* Explanation bubble unit tests

* Merged with main

* updated snapshot

---------

Co-authored-by: Predrag Klepic <v-prklepic@microsoft.com>

src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts

@@ -21,6 +21,9 @@ export const SendQueryRequest = async ({
 }): Promise<void> => {
   if (userPrompt.trim() !== "") {
     useQueryCopilot.getState().setIsGeneratingQuery(true);
+    useQueryCopilot.getState().setShouldIncludeInMessages(true);
+    useQueryCopilot.getState().setShowQueryExplanation(false);
+    useQueryCopilot.getState().setShowExplanationBubble(false);
     useTabs.getState().setIsTabExecuting(true);
     useTabs.getState().setIsQueryErrorThrown(false);
     useQueryCopilot
@@ -62,14 +65,17 @@ export const SendQueryRequest = async ({
         if (generateSQLQueryResponse?.sql) {
           let query = `Here is a query which will help you with provided prompt.\r\n **Prompt:** ${userPrompt}`;
           query += `\r\n${generateSQLQueryResponse.sql}`;
-          useQueryCopilot
+          if (useQueryCopilot.getState().shouldIncludeInMessages) {
-            .getState()
+            useQueryCopilot
-            .setChatMessages([
+              .getState()
-              ...useQueryCopilot.getState().chatMessages,
+              .setChatMessages([
-              { source: 1, message: query, explanation: generateSQLQueryResponse.explanation },
+                ...useQueryCopilot.getState().chatMessages,
-            ]);
+                { source: 1, message: query, explanation: generateSQLQueryResponse.explanation },
-          useQueryCopilot.getState().setGeneratedQuery(generateSQLQueryResponse.sql);
+              ]);
-          useQueryCopilot.getState().setGeneratedQueryComments(generateSQLQueryResponse.explanation);
+            useQueryCopilot.getState().setShowExplanationBubble(true);
+            useQueryCopilot.getState().setGeneratedQuery(generateSQLQueryResponse.sql);
+            useQueryCopilot.getState().setGeneratedQueryComments(generateSQLQueryResponse.explanation);
+          }
         }
       } else {
         handleError(JSON.stringify(generateSQLQueryResponse), "copilotInternalServerError");

src/Explorer/QueryCopilot/V2/Bubbles/Explanation/ExplanationBubble.test.tsx

@@ -0,0 +1,69 @@
+import { Text } from "@fluentui/react";
+import { shallow } from "enzyme";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
+import React from "react";
+import { ExplanationBubble } from "./ExplanationBubble";
+
+describe("Explanation Bubble", () => {
+  const initialStoreState = useQueryCopilot.getState();
+  beforeEach(() => {
+    useQueryCopilot.setState(initialStoreState, true);
+    useQueryCopilot.getState().showExplanationBubble = true;
+    useQueryCopilot.getState().shouldIncludeInMessages = false;
+    jest.useFakeTimers();
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+    jest.useRealTimers();
+  });
+
+  it("should render explanation bubble with generated comments", () => {
+    useQueryCopilot.getState().showQueryExplanation = true;
+    useQueryCopilot.getState().shouldIncludeInMessages = true;
+
+    const wrapper = shallow(<ExplanationBubble />);
+
+    expect(wrapper.find("Stack")).toHaveLength(1);
+    expect(wrapper.find("Text")).toHaveLength(0);
+    expect(wrapper).toMatchSnapshot();
+  });
+
+  it("should render 'Explain this query' link", () => {
+    const mockSetChatMessages = jest.fn();
+    const mockSetIsGeneratingExplanation = jest.fn();
+    const mockSetShouldIncludeInMessages = jest.fn();
+    const mockSetShowQueryExplanation = jest.fn();
+    useQueryCopilot.getState().setChatMessages = mockSetChatMessages;
+    useQueryCopilot.getState().setIsGeneratingExplanation = mockSetIsGeneratingExplanation;
+    useQueryCopilot.getState().setShouldIncludeInMessages = mockSetShouldIncludeInMessages;
+    useQueryCopilot.getState().setShowQueryExplanation = mockSetShowQueryExplanation;
+
+    const wrapper = shallow(<ExplanationBubble />);
+
+    const textElement = wrapper.find(Text);
+    textElement.simulate("click");
+
+    expect(mockSetChatMessages).toHaveBeenCalledWith([
+      ...initialStoreState.chatMessages,
+      { source: 0, message: "Explain this query to me" },
+    ]);
+    expect(mockSetIsGeneratingExplanation).toHaveBeenCalledWith(true);
+    expect(mockSetShouldIncludeInMessages).toHaveBeenCalledWith(true);
+    expect(mockSetShowQueryExplanation).not.toHaveBeenCalled();
+
+    jest.advanceTimersByTime(3000);
+
+    expect(mockSetIsGeneratingExplanation).toHaveBeenCalledWith(false);
+    expect(mockSetShowQueryExplanation).toHaveBeenCalledWith(true);
+  });
+
+  it("should render nothing when conditions are not met", () => {
+    useQueryCopilot.getState().showExplanationBubble = false;
+
+    const wrapper = shallow(<ExplanationBubble />);
+
+    expect(wrapper.isEmptyRender()).toBe(true);
+    expect(wrapper).toMatchSnapshot();
+  });
+});

src/Explorer/QueryCopilot/V2/Bubbles/Explanation/ExplanationBubble.tsx

@@ -0,0 +1,74 @@
+import { Stack, Text } from "@fluentui/react";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
+import React from "react";
+
+export const ExplanationBubble: React.FC = (): JSX.Element => {
+  const {
+    showExplanationBubble,
+    isGeneratingQuery,
+    showQueryExplanation,
+    setShowQueryExplanation,
+    chatMessages,
+    setChatMessages,
+    generatedQueryComments,
+    isGeneratingExplanation,
+    setIsGeneratingExplanation,
+    shouldIncludeInMessages,
+    setShouldIncludeInMessages,
+  } = useQueryCopilot();
+
+  const showExplanation = () => {
+    setChatMessages([...chatMessages, { source: 0, message: "Explain this query to me" }]);
+    setIsGeneratingExplanation(true);
+    setShouldIncludeInMessages(true);
+
+    setTimeout(() => {
+      setIsGeneratingExplanation(false);
+      setShowQueryExplanation(true);
+    }, 3000);
+  };
+
+  return (
+    showExplanationBubble &&
+    !isGeneratingQuery &&
+    !isGeneratingExplanation &&
+    (showQueryExplanation && shouldIncludeInMessages ? (
+      <Stack
+        horizontalAlign="center"
+        tokens={{ padding: 8, childrenGap: 8 }}
+        style={{
+          backgroundColor: "white",
+          borderRadius: "8px",
+          margin: "5px 10px",
+          textAlign: "start",
+        }}
+      >
+        {generatedQueryComments}
+      </Stack>
+    ) : (
+      <Stack
+        style={{
+          display: "flex",
+          alignItems: "center",
+          padding: "5px 5px 5px 40px",
+          margin: "5px",
+          width: "100%",
+        }}
+      >
+        <Text
+          onClick={showExplanation}
+          style={{
+            cursor: "pointer",
+            border: "1.5px solid #B0BEFF",
+            width: "100%",
+            padding: "2px",
+            borderRadius: "4px",
+            marginBottom: "5px",
+          }}
+        >
+          Explain this query to me
+        </Text>
+      </Stack>
+    ))
+  );
+};

src/Explorer/QueryCopilot/V2/Bubbles/Explanation/__snapshots__/ExplanationBubble.test.tsx.snap

@@ -0,0 +1,23 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`Explanation Bubble should render explanation bubble with generated comments 1`] = `
+<Stack
+  horizontalAlign="center"
+  style={
+    Object {
+      "backgroundColor": "white",
+      "borderRadius": "8px",
+      "margin": "5px 10px",
+      "textAlign": "start",
+    }
+  }
+  tokens={
+    Object {
+      "childrenGap": 8,
+      "padding": 8,
+    }
+  }
+/>
+`;
+
+exports[`Explanation Bubble should render nothing when conditions are not met 1`] = `""`;

src/Explorer/QueryCopilot/V2/Sidebar/QueryCopilotSidebar.tsx

@@ -1,5 +1,6 @@
 import { Stack } from "@fluentui/react";
 import { QueryCopilotProps } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
+import { ExplanationBubble } from "Explorer/QueryCopilot/V2/Bubbles/Explanation/ExplanationBubble";
 import { OutputBubble } from "Explorer/QueryCopilot/V2/Bubbles/Output/OutputBubble";
 import { RetrievingBubble } from "Explorer/QueryCopilot/V2/Bubbles/Retriveing/RetrievingBubble";
 import { SampleBubble } from "Explorer/QueryCopilot/V2/Bubbles/Sample/SampleBubble";
@@ -11,7 +12,13 @@ import React from "react";
 import { WelcomeSidebarModal } from "../Modal/WelcomeSidebarModal";
 
 export const QueryCopilotSidebar: React.FC<QueryCopilotProps> = ({ explorer }: QueryCopilotProps): JSX.Element => {
-  const { setWasCopilotUsed, showCopilotSidebar, chatMessages, isGeneratingQuery } = useQueryCopilot();
+  const {
+    setWasCopilotUsed,
+    showCopilotSidebar,
+    chatMessages,
+    isGeneratingQuery,
+    showWelcomeSidebar,
+  } = useQueryCopilot();
 
   React.useEffect(() => {
     if (showCopilotSidebar) {
@@ -22,36 +29,60 @@ export const QueryCopilotSidebar: React.FC<QueryCopilotProps> = ({ explorer }: Q
   return (
     <Stack style={{ width: "100%", height: "100%", backgroundColor: "#FAFAFA" }}>
       <Header />
-      <WelcomeSidebarModal />
+      {showWelcomeSidebar ? (
-      <Stack
+        <WelcomeSidebarModal />
-        style={{
+      ) : (
-          flexGrow: 1,
+        <>
-          display: "flex",
-          flexDirection: "column",
-          overflowY: "auto",
-        }}
-      >
-        <WelcomeBubble />
-        {chatMessages.map((message, index) => (
           <Stack
-            key={index}
-            horizontalAlign="center"
-            tokens={{ padding: 8, childrenGap: 8 }}
             style={{
-              backgroundColor: "#E0E7FF",
+              flexGrow: 1,
-              borderRadius: "8px",
+              display: "flex",
-              margin: "5px 10px",
+              flexDirection: "column",
-              textAlign: "start",
+              overflowY: "auto",
             }}
           >
-            {message}
+            <WelcomeBubble />
+            {chatMessages.map((message, index) =>
+              message.source === 0 ? (
+                <Stack
+                  key={index}
+                  horizontalAlign="center"
+                  tokens={{ padding: 8, childrenGap: 8 }}
+                  style={{
+                    backgroundColor: "#E0E7FF",
+                    borderRadius: "8px",
+                    margin: "5px 10px",
+                    textAlign: "start",
+                  }}
+                >
+                  {message.message}
+                </Stack>
+              ) : (
+                // This part should be wired with OutputBubble
+                <Stack
+                  key={index}
+                  horizontalAlign="center"
+                  tokens={{ padding: 8, childrenGap: 8 }}
+                  style={{
+                    backgroundColor: "white",
+                    borderRadius: "8px",
+                    margin: "5px 10px",
+                    textAlign: "start",
+                  }}
+                >
+                  {message.message}
+                </Stack>
+              )
+            )}
+            <OutputBubble />
+            <RetrievingBubble />
+            <ExplanationBubble />
+
+            {chatMessages.length === 0 && !isGeneratingQuery && <SampleBubble />}
           </Stack>
-        ))}
+          <Footer explorer={explorer} />
-        <OutputBubble />
+        </>
-        <RetrievingBubble />
+      )}
-        {chatMessages.length === 0 && !isGeneratingQuery && <SampleBubble />}
-      </Stack>
-      <Footer explorer={explorer} />
     </Stack>
   );
 };

src/Explorer/QueryCopilot/V2/Sidebar/__snapshots__/QueryCopilotSidebar.test.tsx.snap

@@ -12,51 +12,6 @@ exports[`Query Copilot Sidebar snapshot test should render and not set copilot u
 >
   <Header />
   <WelcomeSidebarModal />
-  <Stack
-    style={
-      Object {
-        "display": "flex",
-        "flexDirection": "column",
-        "flexGrow": 1,
-        "overflowY": "auto",
-      }
-    }
-  >
-    <WelcomeBubble />
-    <OutputBubble />
-    <RetrievingBubble />
-    <SampleBubble />
-  </Stack>
-  <Footer
-    explorer={
-      Explorer {
-        "_isInitializingNotebooks": false,
-        "_resetNotebookWorkspace": [Function],
-        "isFixedCollectionWithSharedThroughputSupported": [Function],
-        "isTabsContentExpanded": [Function],
-        "onRefreshDatabasesKeyPress": [Function],
-        "onRefreshResourcesClick": [Function],
-        "phoenixClient": PhoenixClient {
-          "armResourceId": undefined,
-          "retryOptions": Object {
-            "maxTimeout": 5000,
-            "minTimeout": 5000,
-            "retries": 3,
-          },
-        },
-        "provideFeedbackEmail": [Function],
-        "queriesClient": QueriesClient {
-          "container": [Circular],
-        },
-        "refreshNotebookList": [Function],
-        "resourceTree": ResourceTreeAdapter {
-          "container": [Circular],
-          "copyNotebook": [Function],
-          "parameters": [Function],
-        },
-      }
-    }
-  />
 </Stack>
 `;
 
@@ -72,51 +27,6 @@ exports[`Query Copilot Sidebar snapshot test should render and set copilot used
 >
   <Header />
   <WelcomeSidebarModal />
-  <Stack
-    style={
-      Object {
-        "display": "flex",
-        "flexDirection": "column",
-        "flexGrow": 1,
-        "overflowY": "auto",
-      }
-    }
-  >
-    <WelcomeBubble />
-    <OutputBubble />
-    <RetrievingBubble />
-    <SampleBubble />
-  </Stack>
-  <Footer
-    explorer={
-      Explorer {
-        "_isInitializingNotebooks": false,
-        "_resetNotebookWorkspace": [Function],
-        "isFixedCollectionWithSharedThroughputSupported": [Function],
-        "isTabsContentExpanded": [Function],
-        "onRefreshDatabasesKeyPress": [Function],
-        "onRefreshResourcesClick": [Function],
-        "phoenixClient": PhoenixClient {
-          "armResourceId": undefined,
-          "retryOptions": Object {
-            "maxTimeout": 5000,
-            "minTimeout": 5000,
-            "retries": 3,
-          },
-        },
-        "provideFeedbackEmail": [Function],
-        "queriesClient": QueriesClient {
-          "container": [Circular],
-        },
-        "refreshNotebookList": [Function],
-        "resourceTree": ResourceTreeAdapter {
-          "container": [Circular],
-          "copyNotebook": [Function],
-          "parameters": [Function],
-        },
-      }
-    }
-  />
 </Stack>
 `;
 
@@ -132,51 +42,6 @@ exports[`Query Copilot Sidebar snapshot test should render samples without messa
 >
   <Header />
   <WelcomeSidebarModal />
-  <Stack
-    style={
-      Object {
-        "display": "flex",
-        "flexDirection": "column",
-        "flexGrow": 1,
-        "overflowY": "auto",
-      }
-    }
-  >
-    <WelcomeBubble />
-    <OutputBubble />
-    <RetrievingBubble />
-    <SampleBubble />
-  </Stack>
-  <Footer
-    explorer={
-      Explorer {
-        "_isInitializingNotebooks": false,
-        "_resetNotebookWorkspace": [Function],
-        "isFixedCollectionWithSharedThroughputSupported": [Function],
-        "isTabsContentExpanded": [Function],
-        "onRefreshDatabasesKeyPress": [Function],
-        "onRefreshResourcesClick": [Function],
-        "phoenixClient": PhoenixClient {
-          "armResourceId": undefined,
-          "retryOptions": Object {
-            "maxTimeout": 5000,
-            "minTimeout": 5000,
-            "retries": 3,
-          },
-        },
-        "provideFeedbackEmail": [Function],
-        "queriesClient": QueriesClient {
-          "container": [Circular],
-        },
-        "refreshNotebookList": [Function],
-        "resourceTree": ResourceTreeAdapter {
-          "container": [Circular],
-          "copyNotebook": [Function],
-          "parameters": [Function],
-        },
-      }
-    }
-  />
 </Stack>
 `;
 
@@ -192,69 +57,5 @@ exports[`Query Copilot Sidebar snapshot test should render with chat messages 1`
 >
   <Header />
   <WelcomeSidebarModal />
-  <Stack
-    style={
-      Object {
-        "display": "flex",
-        "flexDirection": "column",
-        "flexGrow": 1,
-        "overflowY": "auto",
-      }
-    }
-  >
-    <WelcomeBubble />
-    <Stack
-      horizontalAlign="center"
-      key="0"
-      style={
-        Object {
-          "backgroundColor": "#E0E7FF",
-          "borderRadius": "8px",
-          "margin": "5px 10px",
-          "textAlign": "start",
-        }
-      }
-      tokens={
-        Object {
-          "childrenGap": 8,
-          "padding": 8,
-        }
-      }
-    >
-      <Component />
-    </Stack>
-    <OutputBubble />
-    <RetrievingBubble />
-  </Stack>
-  <Footer
-    explorer={
-      Explorer {
-        "_isInitializingNotebooks": false,
-        "_resetNotebookWorkspace": [Function],
-        "isFixedCollectionWithSharedThroughputSupported": [Function],
-        "isTabsContentExpanded": [Function],
-        "onRefreshDatabasesKeyPress": [Function],
-        "onRefreshResourcesClick": [Function],
-        "phoenixClient": PhoenixClient {
-          "armResourceId": undefined,
-          "retryOptions": Object {
-            "maxTimeout": 5000,
-            "minTimeout": 5000,
-            "retries": 3,
-          },
-        },
-        "provideFeedbackEmail": [Function],
-        "queriesClient": QueriesClient {
-          "container": [Circular],
-        },
-        "refreshNotebookList": [Function],
-        "resourceTree": ResourceTreeAdapter {
-          "container": [Circular],
-          "copyNotebook": [Function],
-          "parameters": [Function],
-        },
-      }
-    }
-  />
 </Stack>
 `;

src/hooks/useQueryCopilot.ts

@@ -34,6 +34,8 @@ export interface QueryCopilotState {
   chatMessages: CopilotMessage[];
   shouldAllocateContainer: boolean;
   shouldIncludeInMessages: boolean;
+  showExplanationBubble: boolean;
+  showQueryExplanation: boolean;
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void;
   closeFeedbackModal: () => void;
@@ -63,9 +65,10 @@ export interface QueryCopilotState {
   setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => void;
   setShowCopilotSidebar: (showCopilotSidebar: boolean) => void;
   setChatMessages: (chatMessages: CopilotMessage[]) => void;
-
   setShouldAllocateContainer: (shouldAllocateContainer: boolean) => void;
   setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => void;
+  setShowExplanationBubble: (showExplanationBubble: boolean) => void;
+  setShowQueryExplanation: (showQueryExplanation: boolean) => void;
 
   resetQueryCopilotStates: () => void;
 }
@@ -102,6 +105,8 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   chatMessages: [],
   shouldAllocateContainer: true,
   shouldIncludeInMessages: true,
+  showExplanationBubble: false,
+  showQueryExplanation: false,
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
     set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
@@ -135,6 +140,8 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   setChatMessages: (chatMessages: CopilotMessage[]) => set({ chatMessages }),
   setShouldAllocateContainer: (shouldAllocateContainer: boolean) => set({ shouldAllocateContainer }),
   setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => set({ shouldIncludeInMessages }),
+  setShowExplanationBubble: (showExplanationBubble: boolean) => set({ showExplanationBubble }),
+  setShowQueryExplanation: (showQueryExplanation: boolean) => set({ showQueryExplanation }),
 
   resetQueryCopilotStates: () => {
     set((state) => ({
@@ -167,6 +174,8 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       chatMessages: [],
       shouldAllocateContainer: true,
       shouldIncludeInMessages: true,
+      showExplanationBubble: false,
+      showQueryExplanation: false,
     }));
   },
 }));