[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>
2024-11-25 15:06:55 +00:00 · 2023-09-06 19:49:27 +02:00 · 2023-09-06 19:49:27 +02:00 · 76408e2f98
commit 76408e2f98
parent 8c2ca4ab8e
19 changed files with 107 additions and 132 deletions
--- a/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts
+++ b/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}`;
-          query += `\r\n${generateSQLQueryResponse.sql}`;
+          const bubbleMessage = `Here is a query which will help you with provided prompt.\r\n **Prompt:** "${userPrompt}"`;
          if (useQueryCopilot.getState().shouldIncludeInMessages) {
-            useQueryCopilot
-              .getState()
-              .setChatMessages([
+            useQueryCopilot.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);
--- a/src/Explorer/QueryCopilot/Shared/QueryCopilotInterfaces.ts
+++ b/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;
 }

--- a/src/Explorer/QueryCopilot/V2/Bubbles/Explanation/ExplanationBubble.test.tsx
+++ b/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", () => {
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Explanation/ExplanationBubble.tsx
+++ b/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 &&
-    (showQueryExplanation && shouldIncludeInMessages ? (
-      <Stack
-        horizontalAlign="center"
-        tokens={{ padding: 8, childrenGap: 8 }}
-        style={{
-          backgroundColor: "white",
-          borderRadius: "8px",
-          margin: "5px 10px",
-          textAlign: "start",
-        }}
-      >
-        {generatedQueryComments}
-      </Stack>
-    ) : (
+    !isGeneratingExplanation && (
      <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>
-    ))
+    )
  );
 };
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Explanation/snapshots/ExplanationBubble.test.tsx.snap
+++ b/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",
-      "borderRadius": "8px",
-      "margin": "5px 10px",
-      "textAlign": "start",
+      "alignItems": "center",
+      "display": "flex",
+      "margin": "5px",
+      "padding": "5px 5px 5px 50px",
    }
  }
-  tokens={
+>
+  <Text
+    onClick={[Function]}
+    style={
      Object {
-      "childrenGap": 8,
-      "padding": 8,
+        "border": "1.5px solid #B0BEFF",
+        "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`] = `""`;
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Copy/CopyButton.test.tsx
+++ b/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", {});
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Copy/CopyButton.tsx
+++ b/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");
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Feedback/FeedbackButtons.test.tsx
+++ b/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");
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Feedback/FeedbackButtons.tsx
+++ b/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 => {
-  const { generatedQuery, userPrompt } = useQueryCopilot();
+export const FeedbackButtons = ({ sqlQuery }: { sqlQuery: string }): JSX.Element => {
+  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)}
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Insert/InsertButton.test.tsx
+++ b/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";
-    useQueryCopilot.getState().setGeneratedQuery(testQuery);
-    const wrapper = shallow(<InsertButton />);
-
-    const button = wrapper.find(ActionButton).first();
-    button.simulate("click");
-
-    expect(useQueryCopilot.getState().query).toEqual(testQuery);
+    const wrapper = shallow(<InsertButton sqlQuery={""} />);
    expect(wrapper).toMatchSnapshot();
  });
 });
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/Insert/InsertButton.tsx
+++ b/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>
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/OutputBubbleButtons.test.tsx
+++ b/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();
  });
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/OutputBubbleButtons.tsx
+++ b/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 />
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/Buttons/snapshots/OutputBubbleButtons.test.tsx.snap
+++ b/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 />
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/OutputBubble.test.tsx
+++ b/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");
-      useQueryCopilot.getState().setGeneratedQueryComments("test comments");
-      const wrapper = shallow(<OutputBubble />);
+      const wrapper = shallow(
+        <OutputBubble
+          copilotMessage={{ message: "testMessage", source: 1, explanation: "testExplanation", sqlQuery: "testSQL" }}
+        />
+      );

      const editor = wrapper.find(EditorReact).first();

--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/OutputBubble.tsx
+++ b/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" }}>
-        {useQueryCopilot.getState()?.generatedQueryComments}
-      </Stack.Item>
+      <Stack.Item style={{ alignSelf: "flex-start", paddingLeft: "2px" }}>{copilotMessage.message}</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" }}>
--- a/src/Explorer/QueryCopilot/V2/Bubbles/Output/snapshots/OutputBubble.test.tsx.snap
+++ b/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
--- a/src/Explorer/QueryCopilot/V2/Sidebar/QueryCopilotSidebar.tsx
+++ b/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
-                <Stack
-                  key={index}
-                  horizontalAlign="center"
-                  tokens={{ padding: 8, childrenGap: 8 }}
-                  style={{
-                    backgroundColor: "white",
-                    borderRadius: "8px",
-                    margin: "5px 10px",
-                    textAlign: "start",
-                  }}
-                >
-                  {message.message}
-                </Stack>
+                <OutputBubble key={index} copilotMessage={message} />
              )
            )}
-            <OutputBubble />
            <RetrievingBubble />
            <ExplanationBubble />

--- a/src/hooks/useQueryCopilot.ts
+++ b/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,