[Query Copilot] Copilot V2 Retrieving bubble (#1579)

* Changing order of the features

* test name changed

* updates snapshots

* Bubble implementation

* Bubble tests implemented

* Correction for CopilotSampleDB implementation

* rollback to previous query tab changes

---------

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

images/StopGenerating.svg

@@ -0,0 +1,3 @@
+<svg width="14" height="13" viewBox="0 0 14 13" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M3 0C1.61929 0 0.5 1.11929 0.5 2.5V9.5C0.5 10.8807 1.61929 12 3 12H4.75716C4.50353 11.6929 4.28261 11.3578 4.09971 11H3C2.17157 11 1.5 10.3284 1.5 9.5V2.5C1.5 1.67157 2.17157 1 3 1H10C10.8284 1 11.5 1.67157 11.5 2.5V3.59971C11.8578 3.78261 12.1929 4.00353 12.5 4.25716V2.5C12.5 1.11929 11.3807 0 10 0H3ZM9 13C11.4853 13 13.5 10.9853 13.5 8.5C13.5 6.01472 11.4853 4 9 4C6.51472 4 4.5 6.01472 4.5 8.5C4.5 10.9853 6.51472 13 9 13ZM10.8536 6.64645C11.0488 6.84171 11.0488 7.15829 10.8536 7.35355L9.70711 8.5L10.8536 9.64645C11.0488 9.84171 11.0488 10.1583 10.8536 10.3536C10.6583 10.5488 10.3417 10.5488 10.1464 10.3536L9 9.20711L7.85355 10.3536C7.65829 10.5488 7.34171 10.5488 7.14645 10.3536C6.95118 10.1583 6.95118 9.84171 7.14645 9.64645L8.29289 8.5L7.14645 7.35355C6.95118 7.15829 6.95118 6.84171 7.14645 6.64645C7.34171 6.45118 7.65829 6.45118 7.85355 6.64645L9 7.79289L10.1464 6.64645C10.3417 6.45118 10.6583 6.45118 10.8536 6.64645Z" fill="#0078D4"/>
+</svg>

src/Explorer/QueryCopilot/V2/Bubbles/Retriveing/RetrievingBubble.css

@@ -0,0 +1,8 @@
+@keyframes loadingAnimation {
+  0% {
+    transform: translateX(-100%);
+  }
+  100% {
+    transform: translateX(100%);
+  }
+}

src/Explorer/QueryCopilot/V2/Bubbles/Retriveing/RetrievingBubble.test.tsx

@@ -0,0 +1,62 @@
+import { DefaultButton } from "@fluentui/react";
+import { shallow } from "enzyme";
+import React from "react";
+import { RetrievingBubble } from "./RetrievingBubble";
+
+const mockUseQueryCopilot = {
+  isGeneratingQuery: false,
+  setIsGeneratingQuery: jest.fn(),
+  isGeneratingExplanation: false,
+  setIsGeneratingExplanation: jest.fn(),
+  shouldIncludeInMessages: true,
+  setShouldIncludeInMessages: jest.fn(),
+};
+
+jest.mock("hooks/useQueryCopilot", () => ({
+  useQueryCopilot: jest.fn(() => mockUseQueryCopilot),
+}));
+
+describe("RetrievingBubble", () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    mockUseQueryCopilot.isGeneratingQuery = false;
+    mockUseQueryCopilot.isGeneratingExplanation = false;
+    mockUseQueryCopilot.setIsGeneratingQuery.mockClear();
+    mockUseQueryCopilot.setIsGeneratingExplanation.mockClear();
+    mockUseQueryCopilot.setShouldIncludeInMessages.mockClear();
+  });
+
+  it("should render properly when isGeneratingQuery is true", () => {
+    mockUseQueryCopilot.isGeneratingQuery = true;
+    const wrapper = shallow(<RetrievingBubble />);
+    expect(wrapper).toMatchSnapshot();
+  });
+
+  it("should render properly when isGeneratingExplanation is true", () => {
+    mockUseQueryCopilot.isGeneratingExplanation = true;
+    const wrapper = shallow(<RetrievingBubble />);
+    expect(wrapper).toMatchSnapshot();
+  });
+
+  it("when isGeneratingQuery is true clicking stop generating button invokes the correct callbacks", () => {
+    mockUseQueryCopilot.isGeneratingQuery = true;
+
+    const wrapper = shallow(<RetrievingBubble />);
+    wrapper.find(DefaultButton).at(0).simulate("click");
+
+    expect(mockUseQueryCopilot.setIsGeneratingQuery).toHaveBeenCalledWith(false);
+    expect(mockUseQueryCopilot.setIsGeneratingExplanation).toHaveBeenCalledTimes(0);
+    expect(mockUseQueryCopilot.setShouldIncludeInMessages).toHaveBeenCalledWith(false);
+  });
+
+  it("when isGeneratingExplanation is true clicking stop generating button invokes the correct callbacks", () => {
+    mockUseQueryCopilot.isGeneratingExplanation = true;
+
+    const wrapper = shallow(<RetrievingBubble />);
+    wrapper.find(DefaultButton).at(0).simulate("click");
+
+    expect(mockUseQueryCopilot.setIsGeneratingQuery).toHaveBeenCalledTimes(0);
+    expect(mockUseQueryCopilot.setIsGeneratingExplanation).toHaveBeenCalledWith(false);
+    expect(mockUseQueryCopilot.setShouldIncludeInMessages).toHaveBeenCalledWith(false);
+  });
+});

src/Explorer/QueryCopilot/V2/Bubbles/Retriveing/RetrievingBubble.tsx

@@ -0,0 +1,98 @@
+import { DefaultButton, Image, Stack, Text } from "@fluentui/react";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
+import React from "react";
+import StopGeneratingIcon from "../../../../../../images/StopGenerating.svg";
+import "./RetrievingBubble.css";
+
+export const RetrievingBubble = (): JSX.Element => {
+  const {
+    isGeneratingQuery,
+    setIsGeneratingQuery,
+    isGeneratingExplanation,
+    setIsGeneratingExplanation,
+    shouldIncludeInMessages,
+    setShouldIncludeInMessages,
+  } = useQueryCopilot();
+
+  const stopGenerating = () => {
+    if (isGeneratingQuery) {
+      setIsGeneratingQuery(false);
+    }
+    if (isGeneratingExplanation) {
+      setIsGeneratingExplanation(false);
+    }
+    if (shouldIncludeInMessages) {
+      setShouldIncludeInMessages(false);
+    }
+  };
+
+  const bubbleContent = (bubbleType: string) => {
+    return (
+      <Stack
+        horizontalAlign="end"
+        verticalAlign="end"
+        style={{
+          display: "flex",
+          alignItems: "center",
+          padding: "10px",
+          margin: "10px",
+          backgroundColor: "#FAFAFA",
+          borderRadius: "8px",
+        }}
+      >
+        <Text
+          style={{
+            width: "100%",
+            height: "46px",
+            backgroundColor: "white",
+            padding: "12px 16px 16px 16px",
+            gap: "12px",
+            borderRadius: "8px",
+            fontWeight: "bold",
+          }}
+        >
+          Retriveing {bubbleType}
+        </Text>
+        <div
+          style={{
+            width: "100%",
+            height: "4px",
+            backgroundColor: "#E6E6E6",
+            borderRadius: "4px",
+            overflow: "hidden",
+          }}
+        >
+          <div
+            style={{
+              width: "50%",
+              height: "100%",
+              backgroundColor: "#0078D4",
+              animation: "loadingAnimation 2s linear infinite",
+            }}
+          ></div>
+        </div>
+        <Stack
+          horizontalAlign="center"
+          verticalAlign="center"
+          style={{ marginTop: "8px", gap: "8px", alignItems: "center" }}
+        >
+          <DefaultButton
+            onClick={stopGenerating}
+            styles={{ root: { border: "none", background: "none", padding: 0, color: "#424242" } }}
+            style={{ color: "#424242" }}
+            onRenderIcon={() => <Image src={StopGeneratingIcon} />}
+          >
+            Stop generating
+          </DefaultButton>
+        </Stack>
+      </Stack>
+    );
+  };
+
+  return (
+    <>
+      {isGeneratingQuery && bubbleContent("queries")}
+      {isGeneratingExplanation && bubbleContent("explanation")}
+    </>
+  );
+};

src/Explorer/QueryCopilot/V2/Bubbles/Retriveing/__snapshots__/RetrievingBubble.test.tsx.snap

@@ -0,0 +1,183 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`RetrievingBubble should render properly when isGeneratingExplanation is true 1`] = `
+<Fragment>
+  <Stack
+    horizontalAlign="end"
+    style={
+      Object {
+        "alignItems": "center",
+        "backgroundColor": "#FAFAFA",
+        "borderRadius": "8px",
+        "display": "flex",
+        "margin": "10px",
+        "padding": "10px",
+      }
+    }
+    verticalAlign="end"
+  >
+    <Text
+      style={
+        Object {
+          "backgroundColor": "white",
+          "borderRadius": "8px",
+          "fontWeight": "bold",
+          "gap": "12px",
+          "height": "46px",
+          "padding": "12px 16px 16px 16px",
+          "width": "100%",
+        }
+      }
+    >
+      Retriveing 
+      explanation
+    </Text>
+    <div
+      style={
+        Object {
+          "backgroundColor": "#E6E6E6",
+          "borderRadius": "4px",
+          "height": "4px",
+          "overflow": "hidden",
+          "width": "100%",
+        }
+      }
+    >
+      <div
+        style={
+          Object {
+            "animation": "loadingAnimation 2s linear infinite",
+            "backgroundColor": "#0078D4",
+            "height": "100%",
+            "width": "50%",
+          }
+        }
+      />
+    </div>
+    <Stack
+      horizontalAlign="center"
+      style={
+        Object {
+          "alignItems": "center",
+          "gap": "8px",
+          "marginTop": "8px",
+        }
+      }
+      verticalAlign="center"
+    >
+      <CustomizedDefaultButton
+        onClick={[Function]}
+        onRenderIcon={[Function]}
+        style={
+          Object {
+            "color": "#424242",
+          }
+        }
+        styles={
+          Object {
+            "root": Object {
+              "background": "none",
+              "border": "none",
+              "color": "#424242",
+              "padding": 0,
+            },
+          }
+        }
+      >
+        Stop generating
+      </CustomizedDefaultButton>
+    </Stack>
+  </Stack>
+</Fragment>
+`;
+
+exports[`RetrievingBubble should render properly when isGeneratingQuery is true 1`] = `
+<Fragment>
+  <Stack
+    horizontalAlign="end"
+    style={
+      Object {
+        "alignItems": "center",
+        "backgroundColor": "#FAFAFA",
+        "borderRadius": "8px",
+        "display": "flex",
+        "margin": "10px",
+        "padding": "10px",
+      }
+    }
+    verticalAlign="end"
+  >
+    <Text
+      style={
+        Object {
+          "backgroundColor": "white",
+          "borderRadius": "8px",
+          "fontWeight": "bold",
+          "gap": "12px",
+          "height": "46px",
+          "padding": "12px 16px 16px 16px",
+          "width": "100%",
+        }
+      }
+    >
+      Retriveing 
+      queries
+    </Text>
+    <div
+      style={
+        Object {
+          "backgroundColor": "#E6E6E6",
+          "borderRadius": "4px",
+          "height": "4px",
+          "overflow": "hidden",
+          "width": "100%",
+        }
+      }
+    >
+      <div
+        style={
+          Object {
+            "animation": "loadingAnimation 2s linear infinite",
+            "backgroundColor": "#0078D4",
+            "height": "100%",
+            "width": "50%",
+          }
+        }
+      />
+    </div>
+    <Stack
+      horizontalAlign="center"
+      style={
+        Object {
+          "alignItems": "center",
+          "gap": "8px",
+          "marginTop": "8px",
+        }
+      }
+      verticalAlign="center"
+    >
+      <CustomizedDefaultButton
+        onClick={[Function]}
+        onRenderIcon={[Function]}
+        style={
+          Object {
+            "color": "#424242",
+          }
+        }
+        styles={
+          Object {
+            "root": Object {
+              "background": "none",
+              "border": "none",
+              "color": "#424242",
+              "padding": 0,
+            },
+          }
+        }
+      >
+        Stop generating
+      </CustomizedDefaultButton>
+    </Stack>
+  </Stack>
+</Fragment>
+`;

src/Explorer/QueryCopilot/V2/Footer/Footer.tsx

@@ -12,6 +12,7 @@ export const Footer: React.FC = (): JSX.Element => {
     setChatMessages,
     isSamplePromptsOpen,
     setIsSamplePromptsOpen,
+    setIsGeneratingQuery,
   } = useQueryCopilot();
 
   const promptStyles: IButtonStyles = {
@@ -36,6 +37,7 @@ export const Footer: React.FC = (): JSX.Element => {
     if (userPrompt.trim() !== "") {
       setChatMessages([...chatMessages, userPrompt]);
       setUserPrompt("");
+      setIsGeneratingQuery(true);
     }
   };
 

src/Explorer/QueryCopilot/V2/Modal/WelcomeSidebarModal.test.tsx

@@ -2,7 +2,7 @@ import { shallow } from "enzyme";
 import React from "react";
 import { WelcomeSidebarModal } from "./WelcomeSidebarModal";
 
-describe("Footer snapshot test", () => {
+describe("WelcomeSidebarModal snapshot test", () => {
   it("should render ", () => {
     const wrapper = shallow(<WelcomeSidebarModal />);
     expect(wrapper).toMatchSnapshot();

src/Explorer/QueryCopilot/V2/Modal/__snapshots__/WelcomeSidebarModal.test.tsx.snap

@@ -1,6 +1,6 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`Footer snapshot test should render  1`] = `
+exports[`WelcomeSidebarModal snapshot test should render  1`] = `
 <Stack
   style={
     Object {

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

@@ -2,7 +2,7 @@ import { shallow } from "enzyme";
 import React from "react";
 import { QueryCopilotSidebar } from "./QueryCopilotSidebar";
 
-describe("Footer snapshot test", () => {
+describe("QueryCopilotSidebar snapshot test", () => {
   it("should render ", () => {
     const wrapper = shallow(<QueryCopilotSidebar />);
     expect(wrapper).toMatchSnapshot();

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

@@ -1,4 +1,5 @@
 import { Stack } from "@fluentui/react";
+import { RetrievingBubble } from "Explorer/QueryCopilot/V2/Bubbles/Retriveing/RetrievingBubble";
 import { SampleBubble } from "Explorer/QueryCopilot/V2/Bubbles/Sample/SampleBubble";
 import { WelcomeBubble } from "Explorer/QueryCopilot/V2/Bubbles/Welcome/WelcomeBubble";
 import { Footer } from "Explorer/QueryCopilot/V2/Footer/Footer";
@@ -47,6 +48,9 @@ export const QueryCopilotSidebar: React.FC = (): JSX.Element => {
                 {message}
               </Stack>
             ))}
+
+            <RetrievingBubble />
+
             {chatMessages.length === 0 && <SampleBubble />}
           </Stack>
           <Footer />

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

@@ -1,6 +1,6 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`Footer snapshot test should render  1`] = `
+exports[`QueryCopilotSidebar snapshot test should render  1`] = `
 <Stack
   style={
     Object {

src/Platform/Hosted/extractFeatures.ts

@@ -106,9 +106,9 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     enableLegacyMongoShellV2: "true" === get("enablelegacymongoshellv2"),
     enableLegacyMongoShellV2Debug: "true" === get("enablelegacymongoshellv2debug"),
     loadLegacyMongoShellFromBE: "true" === get("loadlegacymongoshellfrombe"),
-    enableCopilot: "true" === get("enablecopilot"),
     enablePriorityBasedThrottling: "true" === get("enableprioritybasedthrottling"),
     enableNPSSurvey: "true" === get("enablenpssurvey"),
+    enableCopilot: "true" === get("enablecopilot"),
     copilotVersion: get("copilotVersion") ? get("copilotVersion") : "v1.0",
   };
 }

src/hooks/useQueryCopilot.ts

@@ -13,6 +13,7 @@ export interface QueryCopilotState {
   query: string;
   selectedQuery: string;
   isGeneratingQuery: boolean;
+  isGeneratingExplanation: boolean;
   isExecuting: boolean;
   dislikeQuery: boolean | undefined;
   showCallout: boolean;
@@ -31,6 +32,7 @@ export interface QueryCopilotState {
   showCopilotSidebar: boolean;
   chatMessages: string[];
   shouldAllocateContainer: boolean;
+  shouldIncludeInMessages: boolean;
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void;
   closeFeedbackModal: () => void;
@@ -41,6 +43,7 @@ export interface QueryCopilotState {
   setGeneratedQuery: (generatedQuery: string) => void;
   setSelectedQuery: (selectedQuery: string) => void;
   setIsGeneratingQuery: (isGeneratingQuery: boolean) => void;
+  setIsGeneratingExplanation: (isGeneratingExplanation: boolean) => void;
   setIsExecuting: (isExecuting: boolean) => void;
   setLikeQuery: (likeQuery: boolean) => void;
   setDislikeQuery: (dislikeQuery: boolean | undefined) => void;
@@ -61,6 +64,7 @@ export interface QueryCopilotState {
   setChatMessages: (chatMessages: string[]) => void;
 
   setShouldAllocateContainer: (shouldAllocateContainer: boolean) => void;
+  setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => void;
 
   resetQueryCopilotStates: () => void;
 }
@@ -77,6 +81,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   query: "",
   selectedQuery: "",
   isGeneratingQuery: false,
+  isGeneratingExplanation: false,
   isExecuting: false,
   dislikeQuery: undefined,
   showCallout: false,
@@ -95,6 +100,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   showCopilotSidebar: false,
   chatMessages: [],
   shouldAllocateContainer: true,
+  shouldIncludeInMessages: true,
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
     set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
@@ -107,6 +113,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   setGeneratedQuery: (generatedQuery: string) => set({ generatedQuery }),
   setSelectedQuery: (selectedQuery: string) => set({ selectedQuery }),
   setIsGeneratingQuery: (isGeneratingQuery: boolean) => set({ isGeneratingQuery }),
+  setIsGeneratingExplanation: (isGeneratingExplanation: boolean) => set({ isGeneratingExplanation }),
   setIsExecuting: (isExecuting: boolean) => set({ isExecuting }),
   setLikeQuery: (likeQuery: boolean) => set({ likeQuery }),
   setDislikeQuery: (dislikeQuery: boolean | undefined) => set({ dislikeQuery }),
@@ -126,6 +133,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   setShowCopilotSidebar: (showCopilotSidebar: boolean) => set({ showCopilotSidebar }),
   setChatMessages: (chatMessages: string[]) => set({ chatMessages }),
   setShouldAllocateContainer: (shouldAllocateContainer: boolean) => set({ shouldAllocateContainer }),
+  setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => set({ shouldIncludeInMessages }),
 
   resetQueryCopilotStates: () => {
     set((state) => ({
@@ -139,6 +147,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       query: "",
       selectedQuery: "",
       isGeneratingQuery: false,
+      isGeneratingExplanation: false,
       isExecuting: false,
       dislikeQuery: undefined,
       showCallout: false,
@@ -156,6 +165,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       showCopilotSidebar: false,
       chatMessages: [],
       shouldAllocateContainer: true,
+      shouldIncludeInMessages: true,
     }));
   },
 }));