[Query Copilot V2] Wire and adjust Output bubble with backend communication (#1599)

* Initial wiring of copilot backend and bubble

* Additional changes in explanation bubbles

* Changes based on checks

* test snapshots updated

---------

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

src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts

@@ -20,15 +20,13 @@ export const SendQueryRequest = async ({
   explorer: Explorer;
 }): 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
       .getState()
       .setChatMessages([...useQueryCopilot.getState().chatMessages, { source: 0, message: userPrompt }]);
+    useQueryCopilot.getState().setIsGeneratingQuery(true);
+    useQueryCopilot.getState().setShouldIncludeInMessages(true);
+    useTabs.getState().setIsTabExecuting(true);
+    useTabs.getState().setIsQueryErrorThrown(false);
     try {
       if (
         useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
@@ -62,14 +60,16 @@ export const SendQueryRequest = async ({
       const generateSQLQueryResponse: GenerateSQLQueryResponse = await response?.json();
       if (response.ok) {
         if (generateSQLQueryResponse?.sql) {
-          let query = `Here is a query which will help you with provided prompt.\r\n **Prompt:** ${userPrompt}`;
+          const bubbleMessage = `Here is a query which will help you with provided prompt.\r\n **Prompt:** "${userPrompt}"`;
-          query += `\r\n${generateSQLQueryResponse.sql}`;
           if (useQueryCopilot.getState().shouldIncludeInMessages) {
-            useQueryCopilot
+            useQueryCopilot.getState().setChatMessages([
-              .getState()
-              .setChatMessages([
               ...useQueryCopilot.getState().chatMessages,
-                { source: 1, message: query, explanation: generateSQLQueryResponse.explanation },
+              {
+                source: 1,
+                message: bubbleMessage,
+                sqlQuery: generateSQLQueryResponse.sql,
+                explanation: generateSQLQueryResponse.explanation,
+              },
             ]);
             useQueryCopilot.getState().setShowExplanationBubble(true);
             useQueryCopilot.getState().setGeneratedQuery(generateSQLQueryResponse.sql);

src/Explorer/QueryCopilot/Shared/QueryCopilotInterfaces.ts

@@ -11,11 +11,13 @@ export interface GenerateSQLQueryResponse {
 enum MessageSource {
   User,
   AI,
+  AIExplanation,
 }
 
 export interface CopilotMessage {
   source: MessageSource;
   message: string;
+  sqlQuery?: string;
   explanation?: string;
 }
 

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

@@ -19,25 +19,25 @@ describe("Explanation Bubble", () => {
   });
 
   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.find("Text")).toHaveLength(1);
     expect(wrapper).toMatchSnapshot();
   });
 
   it("should render 'Explain this query' link", () => {
+    useQueryCopilot.getState().shouldIncludeInMessages = true;
     const mockSetChatMessages = jest.fn();
     const mockSetIsGeneratingExplanation = jest.fn();
     const mockSetShouldIncludeInMessages = jest.fn();
-    const mockSetShowQueryExplanation = jest.fn();
+    const mockSetShowExplanationBubble = jest.fn();
     useQueryCopilot.getState().setChatMessages = mockSetChatMessages;
     useQueryCopilot.getState().setIsGeneratingExplanation = mockSetIsGeneratingExplanation;
     useQueryCopilot.getState().setShouldIncludeInMessages = mockSetShouldIncludeInMessages;
-    useQueryCopilot.getState().setShowQueryExplanation = mockSetShowQueryExplanation;
+    useQueryCopilot.getState().setShowExplanationBubble = mockSetShowExplanationBubble;
 
     const wrapper = shallow(<ExplanationBubble />);
 
@@ -50,12 +50,12 @@ describe("Explanation Bubble", () => {
     ]);
     expect(mockSetIsGeneratingExplanation).toHaveBeenCalledWith(true);
     expect(mockSetShouldIncludeInMessages).toHaveBeenCalledWith(true);
-    expect(mockSetShowQueryExplanation).not.toHaveBeenCalled();
+    expect(mockSetShowExplanationBubble).toHaveBeenCalledWith(false);
 
     jest.advanceTimersByTime(3000);
 
     expect(mockSetIsGeneratingExplanation).toHaveBeenCalledWith(false);
-    expect(mockSetShowQueryExplanation).toHaveBeenCalledWith(true);
+    expect(mockSetChatMessages).toHaveBeenCalled();
   });
 
   it("should render nothing when conditions are not met", () => {

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

@@ -6,53 +6,39 @@ export const ExplanationBubble: React.FC = (): JSX.Element => {
   const {
     showExplanationBubble,
     isGeneratingQuery,
-    showQueryExplanation,
-    setShowQueryExplanation,
     chatMessages,
     setChatMessages,
     generatedQueryComments,
     isGeneratingExplanation,
     setIsGeneratingExplanation,
-    shouldIncludeInMessages,
     setShouldIncludeInMessages,
+    setShowExplanationBubble,
   } = useQueryCopilot();
 
   const showExplanation = () => {
     setChatMessages([...chatMessages, { source: 0, message: "Explain this query to me" }]);
     setIsGeneratingExplanation(true);
     setShouldIncludeInMessages(true);
+    setShowExplanationBubble(false);
 
     setTimeout(() => {
+      if (useQueryCopilot.getState().shouldIncludeInMessages) {
         setIsGeneratingExplanation(false);
-      setShowQueryExplanation(true);
+        setChatMessages([...chatMessages, { source: 2, message: generatedQueryComments }]);
+      }
     }, 3000);
   };
 
   return (
     showExplanationBubble &&
     !isGeneratingQuery &&
-    !isGeneratingExplanation &&
+    !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",
+          padding: "5px 5px 5px 50px",
           margin: "5px",
-          width: "100%",
         }}
       >
         <Text
@@ -69,6 +55,6 @@ export const ExplanationBubble: React.FC = (): JSX.Element => {
           Explain this query to me
         </Text>
       </Stack>
-    ))
+    )
   );
 };

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

@@ -2,22 +2,31 @@
 
 exports[`Explanation Bubble should render explanation bubble with generated comments 1`] = `
 <Stack
-  horizontalAlign="center"
   style={
     Object {
-      "backgroundColor": "white",
+      "alignItems": "center",
-      "borderRadius": "8px",
+      "display": "flex",
-      "margin": "5px 10px",
+      "margin": "5px",
-      "textAlign": "start",
+      "padding": "5px 5px 5px 50px",
     }
   }
-  tokens={
+>
+  <Text
+    onClick={[Function]}
+    style={
       Object {
-      "childrenGap": 8,
+        "border": "1.5px solid #B0BEFF",
-      "padding": 8,
+        "borderRadius": "4px",
+        "cursor": "pointer",
+        "marginBottom": "5px",
+        "padding": "2px",
+        "width": "100%",
       }
     }
-/>
+  >
+    Explain this query to me
+  </Text>
+</Stack>
 `;
 
 exports[`Explanation Bubble should render nothing when conditions are not met 1`] = `""`;

src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Copy/CopyButton.test.tsx

@@ -10,7 +10,7 @@ describe("Copy button snapshot tests", () => {
   it("should render and click copy", async () => {
     const testInput = "test input query";
     useQueryCopilot.getState().setGeneratedQuery(testInput);
-    const wrapper = shallow(<CopyButton />);
+    const wrapper = shallow(<CopyButton sqlQuery={""} />);
 
     const button = wrapper.find(IconButton).first();
     button.simulate("click", {});

src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Copy/CopyButton.tsx

@@ -1,12 +1,11 @@
 import { IconButton } from "@fluentui/react";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
 import React from "react";
 import CopilotCopy from "../../../../../../../../images/CopilotCopy.svg";
 
-export const CopyButton: React.FC = (): JSX.Element => {
+export const CopyButton = ({ sqlQuery }: { sqlQuery: string }): JSX.Element => {
   const copyGeneratedCode = (): void => {
     const queryElement = document.createElement("textarea");
-    queryElement.value = useQueryCopilot.getState().generatedQuery;
+    queryElement.value = sqlQuery;
     document.body.appendChild(queryElement);
     queryElement.select();
     document.execCommand("copy");

src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Feedback/FeedbackButtons.test.tsx

@@ -20,7 +20,7 @@ beforeEach(() => {
 
 describe("Feedback buttons snapshot tests", () => {
   it("should click like and show callout", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let likeButton = wrapper.find(IconButton).first();
     const dislikeButton = wrapper.find(IconButton).last();
@@ -35,7 +35,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should click like and dismiss callout", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     const likeButton = wrapper.find(IconButton).first();
     likeButton.simulate("click");
@@ -49,7 +49,7 @@ describe("Feedback buttons snapshot tests", () => {
 
   it("should click like and submit feedback", () => {
     const spy = jest.spyOn(useQueryCopilot.getState(), "openFeedbackModal");
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     const likeButton = wrapper.find(IconButton).first();
     likeButton.simulate("click");
@@ -61,7 +61,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should hover over like", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let likeButton = wrapper.find(IconButton).first();
     likeButton.simulate("mouseover");
@@ -72,7 +72,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should hover over rest like and leave", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let likeButton = wrapper.find(IconButton).first();
     likeButton.simulate("mouseover");
@@ -84,7 +84,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should hover over pressed like and leave", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let likeButton = wrapper.find(IconButton).first();
     likeButton.simulate("click");
@@ -97,7 +97,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should hover over like and click", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let likeButton = wrapper.find(IconButton).first();
     likeButton.simulate("mouseover");
@@ -109,7 +109,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should dobule click on like", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let likeButton = wrapper.find(IconButton).first();
     likeButton.simulate("click");
@@ -124,7 +124,7 @@ describe("Feedback buttons snapshot tests", () => {
 
   it("should click dislike and show popup", () => {
     const spy = jest.spyOn(useQueryCopilot.getState(), "openFeedbackModal");
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     const likeButton = wrapper.find(IconButton).first();
     let dislikeButton = wrapper.find(IconButton).last();
@@ -140,7 +140,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should hover over dislike", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let dislikeButton = wrapper.find(IconButton).last();
     dislikeButton.simulate("mouseover");
@@ -151,7 +151,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should hover over rest dislike and leave", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let dislikeButton = wrapper.find(IconButton).last();
     dislikeButton.simulate("mouseover");
@@ -163,7 +163,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should hover over pressed dislike and leave", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let dislikeButton = wrapper.find(IconButton).last();
     dislikeButton.simulate("click");
@@ -178,7 +178,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should hover over dislike and click", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let dislikeButton = wrapper.find(IconButton).last();
     dislikeButton.simulate("mouseover");
@@ -190,7 +190,7 @@ describe("Feedback buttons snapshot tests", () => {
   });
 
   it("should dobule click on dislike", () => {
-    const wrapper = shallow(<FeedbackButtons />);
+    const wrapper = shallow(<FeedbackButtons sqlQuery={""} />);
 
     let dislikeButton = wrapper.find(IconButton).last();
     dislikeButton.simulate("click");

src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Feedback/FeedbackButtons.tsx

@@ -6,8 +6,8 @@ import LikeHover from "../../../../../../../../images/CopilotLikeHover.svg";
 import LikePressed from "../../../../../../../../images/CopilotLikePressed.svg";
 import LikeRest from "../../../../../../../../images/CopilotLikeRest.svg";
 
-export const FeedbackButtons: React.FC = (): JSX.Element => {
+export const FeedbackButtons = ({ sqlQuery }: { sqlQuery: string }): JSX.Element => {
-  const { generatedQuery, userPrompt } = useQueryCopilot();
+  const { userPrompt } = useQueryCopilot();
 
   const [likeQuery, setLikeQuery] = useState<boolean>(false);
   const [dislikeQuery, setDislikeQuery] = useState<boolean>(false);
@@ -35,7 +35,7 @@ export const FeedbackButtons: React.FC = (): JSX.Element => {
               <Link
                 onClick={() => {
                   setCalloutVisible(false);
-                  useQueryCopilot.getState().openFeedbackModal(generatedQuery, true, userPrompt);
+                  useQueryCopilot.getState().openFeedbackModal(sqlQuery, true, userPrompt);
                 }}
               >
                 more feedback?
@@ -82,7 +82,7 @@ export const FeedbackButtons: React.FC = (): JSX.Element => {
             setLikeQuery(false);
             setDislikeImageLink(LikePressed);
             setLikeImageLink(LikeRest);
-            useQueryCopilot.getState().openFeedbackModal(generatedQuery, false, userPrompt);
+            useQueryCopilot.getState().openFeedbackModal(sqlQuery, false, userPrompt);
           }
         }}
         onMouseOver={() => setDislikeImageLink(LikeHover)}

src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Insert/InsertButton.test.tsx

@@ -1,19 +1,10 @@
-import { ActionButton } from "@fluentui/react";
 import { InsertButton } from "Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Insert/InsertButton";
 import { shallow } from "enzyme";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
 import React from "react";
 
 describe("Insert button snapshot tests", () => {
   it("should click and update state", () => {
-    const testQuery = "test query";
+    const wrapper = shallow(<InsertButton sqlQuery={""} />);
-    useQueryCopilot.getState().setGeneratedQuery(testQuery);
-    const wrapper = shallow(<InsertButton />);
-
-    const button = wrapper.find(ActionButton).first();
-    button.simulate("click");
-
-    expect(useQueryCopilot.getState().query).toEqual(testQuery);
     expect(wrapper).toMatchSnapshot();
   });
 });

src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Insert/InsertButton.tsx

@@ -3,12 +3,12 @@ import { useQueryCopilot } from "hooks/useQueryCopilot";
 import React from "react";
 import CopilotInsert from "../../../../../../../../images/CopilotInsert.svg";
 
-export const InsertButton: React.FC = (): JSX.Element => {
+export const InsertButton = ({ sqlQuery }: { sqlQuery: string }): JSX.Element => {
   return (
     <ActionButton
       iconProps={{ imageProps: { src: CopilotInsert } }}
       style={{ borderRadius: "4px", borderWidth: "1px", borderColor: "#D1D1D1", height: "24px", paddingBottom: "2px" }}
-      onClick={() => useQueryCopilot.getState().setQuery(useQueryCopilot.getState().generatedQuery)}
+      onClick={() => useQueryCopilot.getState().setQuery(sqlQuery)}
     >
       Insert
     </ActionButton>

src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/OutputBubbleButtons.test.tsx

@@ -4,7 +4,7 @@ import React from "react";
 
 describe("Output Bubble Buttons snapshot tests", () => {
   it("should render", () => {
-    const wrapper = shallow(<OutputBubbleButtons />);
+    const wrapper = shallow(<OutputBubbleButtons sqlQuery={""} />);
 
     expect(wrapper).toMatchSnapshot();
   });

src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/OutputBubbleButtons.tsx

@@ -5,17 +5,17 @@ import { InsertButton } from "Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/In
 import { MoreButton } from "Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/More/MoreButton";
 import React from "react";
 
-export const OutputBubbleButtons: React.FC = (): JSX.Element => {
+export const OutputBubbleButtons = ({ sqlQuery }: { sqlQuery: string }): JSX.Element => {
   return (
     <Stack horizontal>
       <Stack.Item style={{ paddingTop: "5px" }}>
-        <InsertButton />
+        <InsertButton sqlQuery={sqlQuery} />
       </Stack.Item>
       <Stack.Item>
-        <CopyButton />
+        <CopyButton sqlQuery={sqlQuery} />
       </Stack.Item>
       <Stack.Item>
-        <FeedbackButtons />
+        <FeedbackButtons sqlQuery={sqlQuery} />
       </Stack.Item>
       <Stack.Item>
         <MoreButton />

src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/__snapshots__/OutputBubbleButtons.test.tsx.snap

@@ -11,13 +11,19 @@ exports[`Output Bubble Buttons snapshot tests should render 1`] = `
       }
     }
   >
-    <InsertButton />
+    <InsertButton
+      sqlQuery=""
+    />
   </StackItem>
   <StackItem>
-    <CopyButton />
+    <CopyButton
+      sqlQuery=""
+    />
   </StackItem>
   <StackItem>
-    <FeedbackButtons />
+    <FeedbackButtons
+      sqlQuery=""
+    />
   </StackItem>
   <StackItem>
     <MoreButton />

src/Explorer/QueryCopilot/V2/Bubbles/Output/OutputBubble.test.tsx

@@ -1,16 +1,17 @@
 import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
 import { OutputBubble } from "Explorer/QueryCopilot/V2/Bubbles/Output/OutputBubble";
 import { shallow } from "enzyme";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { withHooks } from "jest-react-hooks-shallow";
 import React from "react";
 
 describe("Output Bubble snapshot tests", () => {
   it("should render and update height", () => {
     withHooks(() => {
-      useQueryCopilot.getState().setGeneratedQuery("test query");
+      const wrapper = shallow(
-      useQueryCopilot.getState().setGeneratedQueryComments("test comments");
+        <OutputBubble
-      const wrapper = shallow(<OutputBubble />);
+          copilotMessage={{ message: "testMessage", source: 1, explanation: "testExplanation", sqlQuery: "testSQL" }}
+        />
+      );
 
       const editor = wrapper.find(EditorReact).first();
 

src/Explorer/QueryCopilot/V2/Bubbles/Output/OutputBubble.tsx

@@ -1,10 +1,10 @@
 import { Stack, Text } from "@fluentui/react";
 import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
+import { CopilotMessage } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
 import { OutputBubbleButtons } from "Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/OutputBubbleButtons";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
 import React, { useState } from "react";
 
-export const OutputBubble: React.FC = (): JSX.Element => {
+export const OutputBubble = ({ copilotMessage }: { copilotMessage: CopilotMessage }): JSX.Element => {
   const [windowHeight, setWindowHeight] = useState<string>();
 
   const calculateQueryWindowHeight = (): string => {
@@ -29,13 +29,11 @@ export const OutputBubble: React.FC = (): JSX.Element => {
       }}
       tokens={{ padding: 8, childrenGap: 8 }}
     >
-      <Stack.Item style={{ alignSelf: "flex-start", paddingLeft: "2px" }}>
+      <Stack.Item style={{ alignSelf: "flex-start", paddingLeft: "2px" }}>{copilotMessage.message}</Stack.Item>
-        {useQueryCopilot.getState()?.generatedQueryComments}
-      </Stack.Item>
       <Stack.Item style={{ alignSelf: "stretch", flexGrow: 4 }}>
         <EditorReact
           language={"sql"}
-          content={useQueryCopilot.getState()?.generatedQuery}
+          content={copilotMessage.sqlQuery}
           isReadOnly={true}
           ariaLabel={"AI Response"}
           wordWrap="on"
@@ -48,7 +46,7 @@ export const OutputBubble: React.FC = (): JSX.Element => {
         />
       </Stack.Item>
       <Stack.Item style={{ alignSelf: "flex-start" }}>
-        <OutputBubbleButtons />
+        <OutputBubbleButtons sqlQuery={copilotMessage.sqlQuery} />
       </Stack.Item>
       <Stack.Item>
         <Text style={{ fontWeight: 400, fontSize: "10px", lineHeight: "14px" }}>

src/Explorer/QueryCopilot/V2/Bubbles/Output/__snapshots__/OutputBubble.test.tsx.snap

@@ -28,7 +28,7 @@ exports[`Output Bubble snapshot tests should render and update height 1`] = `
       }
     }
   >
-    test comments
+    testMessage
   </StackItem>
   <StackItem
     style={
@@ -40,7 +40,7 @@ exports[`Output Bubble snapshot tests should render and update height 1`] = `
   >
     <EditorReact
       ariaLabel="AI Response"
-      content="test query"
+      content="testSQL"
       isReadOnly={true}
       language="sql"
       lineDecorationsWidth={0}
@@ -68,7 +68,9 @@ exports[`Output Bubble snapshot tests should render and update height 1`] = `
       }
     }
   >
-    <OutputBubbleButtons />
+    <OutputBubbleButtons
+      sqlQuery="testSQL"
+    />
   </StackItem>
   <StackItem>
     <Text

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

@@ -43,13 +43,13 @@ export const QueryCopilotSidebar: React.FC<QueryCopilotProps> = ({ explorer }: Q
           >
             <WelcomeBubble />
             {chatMessages.map((message, index) =>
-              message.source === 0 ? (
+              message.source === 0 || message.source === 2 ? (
                 <Stack
                   key={index}
                   horizontalAlign="center"
                   tokens={{ padding: 8, childrenGap: 8 }}
                   style={{
-                    backgroundColor: "#E0E7FF",
+                    backgroundColor: message.source === 0 ? "#E0E7FF" : "white",
                     borderRadius: "8px",
                     margin: "5px 10px",
                     textAlign: "start",
@@ -58,23 +58,9 @@ export const QueryCopilotSidebar: React.FC<QueryCopilotProps> = ({ explorer }: Q
                   {message.message}
                 </Stack>
               ) : (
-                // This part should be wired with OutputBubble
+                <OutputBubble key={index} copilotMessage={message} />
-                <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 />
 

src/hooks/useQueryCopilot.ts

@@ -37,7 +37,6 @@ export interface QueryCopilotState {
   chatMessages: CopilotMessage[];
   shouldIncludeInMessages: boolean;
   showExplanationBubble: boolean;
-  showQueryExplanation: boolean;
   notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo;
   containerStatus: ContainerInfo;
   isAllocatingContainer: boolean;
@@ -72,7 +71,6 @@ export interface QueryCopilotState {
   setChatMessages: (chatMessages: CopilotMessage[]) => void;
   setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => void;
   setShowExplanationBubble: (showExplanationBubble: boolean) => void;
-  setShowQueryExplanation: (showQueryExplanation: boolean) => void;
   setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => void;
   setContainerStatus: (containerStatus: ContainerInfo) => void;
   setIsAllocatingContainer: (isAllocatingContainer: boolean) => void;
@@ -113,7 +111,6 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   chatMessages: [],
   shouldIncludeInMessages: true,
   showExplanationBubble: false,
-  showQueryExplanation: false,
   notebookServerInfo: {
     notebookServerEndpoint: undefined,
     authToken: undefined,
@@ -158,7 +155,6 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   setChatMessages: (chatMessages: CopilotMessage[]) => set({ chatMessages }),
   setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => set({ shouldIncludeInMessages }),
   setShowExplanationBubble: (showExplanationBubble: boolean) => set({ showExplanationBubble }),
-  setShowQueryExplanation: (showQueryExplanation: boolean) => set({ showQueryExplanation }),
   setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
     set({ notebookServerInfo }),
   setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }),
@@ -206,7 +202,6 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       chatMessages: [],
       shouldIncludeInMessages: true,
       showExplanationBubble: false,
-      showQueryExplanation: false,
       notebookServerInfo: {
         notebookServerEndpoint: undefined,
         authToken: undefined,