[Query Copilot] Version 2 initial code (#1566)

* Query Copilot Sidecar initial commit

* additional improvements with the Welcome pop

* Additional implementation and messages for sidecar

* introducing copilot version

* Renaming from car to bar

* Image names changed

* fixing PR errors

* additional changes related with the versions

* Additional implementations and fixes

* Removing unused interface

* Additional styling changes and state management

* Additional changes related with Sidebar

---------

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

src/Explorer/ContextMenuButtonFactory.tsx

@@ -1,3 +1,4 @@
+import { useDatabases } from "Explorer/useDatabases";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { traceOpen } from "Shared/Telemetry/TelemetryProcessor";
 import { ReactTabKind, useTabs } from "hooks/useTabs";
@@ -146,14 +147,24 @@ export const createCollectionContextMenuButton = (
 export const createSampleCollectionContextMenuButton = (): TreeNodeMenuItem[] => {
   const items: TreeNodeMenuItem[] = [];
   if (userContext.apiType === "SQL") {
-    items.push({
-      iconSrc: AddSqlQueryIcon,
-      onClick: () => {
-        useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot);
-        traceOpen(Action.OpenQueryCopilotFromNewQuery, { apiType: userContext.apiType });
-      },
-      label: "New SQL Query",
-    });
+    const copilotVersion = userContext.features.copilotVersion;
+    if (copilotVersion === "v1.0") {
+      items.push({
+        iconSrc: AddSqlQueryIcon,
+        onClick: () => {
+          useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot);
+          traceOpen(Action.OpenQueryCopilotFromNewQuery, { apiType: userContext.apiType });
+        },
+        label: "New SQL Query",
+      });
+    } else if (copilotVersion === "v2.0") {
+      const sampleCollection = useDatabases.getState().sampleDataResourceTokenCollection;
+      items.push({
+        iconSrc: AddSqlQueryIcon,
+        onClick: () => sampleCollection && sampleCollection.onNewQueryClick(sampleCollection, undefined),
+        label: "New SQL Query",
+      });
+    }
   }
 
   return items;

src/Explorer/QueryCopilot/Sidebar/QueryCopilotSidebar.tsx

@@ -0,0 +1,243 @@
+import { IButtonStyles, IconButton, Image, Stack, Text, TextField } from "@fluentui/react";
+import { SamplePrompts, SamplePromptsProps } from "Explorer/QueryCopilot/SamplePrompts/SamplePrompts";
+import { WelcomeSidebarPopup } from "Explorer/QueryCopilot/Sidebar/WelcomeSidebarPopup";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
+import React from "react";
+import CopilotIcon from "../../../../images/CopilotSidebarLogo.svg";
+import HintIcon from "../../../../images/Hint.svg";
+
+const sampleChatMessages: string[] = [
+  "Write a query to return last 10 records in the database",
+  'Write a query to return all records in this table created in the last thirty days which also have the record owner as "Contoso"',
+];
+
+const promptStyles: IButtonStyles = {
+  root: { border: "5px", selectors: { ":hover": { outline: "1px dashed #605e5c" } } },
+  label: { fontWeight: 400, textAlign: "left", paddingLeft: 8 },
+};
+
+export const QueryCopilotSidebar: React.FC = (): JSX.Element => {
+  const {
+    setWasCopilotUsed,
+    showCopilotSidebar,
+    setShowCopilotSidebar,
+    userPrompt,
+    setUserPrompt,
+    chatMessages,
+    setChatMessages,
+    showWelcomeSidebar,
+    isSamplePromptsOpen,
+    setIsSamplePromptsOpen,
+  } = useQueryCopilot();
+
+  const sampleProps: SamplePromptsProps = {
+    isSamplePromptsOpen: isSamplePromptsOpen,
+    setIsSamplePromptsOpen: setIsSamplePromptsOpen,
+    setTextBox: setUserPrompt,
+  };
+
+  const handleSendMessage = () => {
+    if (userPrompt.trim() !== "") {
+      setChatMessages([...chatMessages, userPrompt]);
+      setUserPrompt("");
+    }
+  };
+
+  const handleInputChange = (value: string) => {
+    setUserPrompt(value);
+  };
+
+  const handleEnterKeyPress = (event: React.KeyboardEvent<HTMLInputElement>) => {
+    if (event.key === "Enter" && !event.shiftKey) {
+      event.preventDefault();
+      handleSendMessage();
+    }
+  };
+
+  React.useEffect(() => {
+    if (showCopilotSidebar) {
+      setWasCopilotUsed(true);
+    }
+  }, []);
+
+  return (
+    <Stack style={{ width: "100%", height: "100%", backgroundColor: "#FAFAFA", overflow: "auto" }}>
+      <Stack style={{ margin: "15px 0px 0px 0px", padding: "5px" }}>
+        <Stack style={{ display: "flex", justifyContent: "space-between" }} horizontal verticalAlign="center">
+          <Stack horizontal verticalAlign="center">
+            <Image src={CopilotIcon} />
+            <Text style={{ marginLeft: "5px", fontWeight: "bold" }}>Copilot</Text>
+            <Text
+              style={{
+                background: "#f0f0f0",
+                fontSize: "10px",
+                padding: "2px 4px",
+                marginLeft: "5px",
+                borderRadius: "8px",
+              }}
+            >
+              Preview
+            </Text>
+          </Stack>
+          <IconButton
+            onClick={() => setShowCopilotSidebar(false)}
+            iconProps={{ iconName: "Cancel" }}
+            title="Exit"
+            ariaLabel="Exit"
+            style={{ color: "#424242", verticalAlign: "middle" }}
+          />
+        </Stack>
+      </Stack>
+      {showWelcomeSidebar ? (
+        <Stack.Item styles={{ root: { textAlign: "center", verticalAlign: "middle" } }}>
+          <WelcomeSidebarPopup />
+        </Stack.Item>
+      ) : (
+        <>
+          <Stack horizontalAlign="center" style={{ color: "#707070" }} tokens={{ padding: 8, childrenGap: 8 }}>
+            {new Date().toLocaleDateString("en-US", {
+              month: "long",
+              day: "numeric",
+            })}{" "}
+            {new Date().toLocaleTimeString("en-US", {
+              hour: "numeric",
+              minute: "numeric",
+              hour12: true,
+            })}
+          </Stack>
+          <Stack style={{ flexGrow: 1, display: "flex", flexDirection: "column" }}>
+            <div style={{ flexGrow: 1, overflowY: "auto" }}>
+              <Stack
+                tokens={{ padding: 16, childrenGap: 12 }}
+                style={{
+                  backgroundColor: "white",
+                  margin: "5px 10px",
+                  borderRadius: "8px",
+                }}
+              >
+                <Text variant="medium">
+                  Hello, I am Cosmos Db&apos;s copilot assistant. I can help you do the following things:
+                </Text>
+                <Stack tokens={{ childrenGap: 8 }}>
+                  <Stack horizontal style={{ marginLeft: "15px" }} tokens={{ childrenGap: 8 }}>
+                    <Text style={{ fontSize: "16px", lineHeight: "16px", verticalAlign: "middle" }}>•</Text>
+                    <Text style={{ verticalAlign: "middle" }}>Generate queries based upon prompt you suggest</Text>
+                  </Stack>
+                  <Stack horizontal style={{ marginLeft: "15px" }} tokens={{ childrenGap: 8 }}>
+                    <Text style={{ fontSize: "16px", lineHeight: "16px", verticalAlign: "middle" }}>•</Text>
+                    <Text style={{ verticalAlign: "middle" }}>
+                      Explain and provide alternate queries for a query suggested by you
+                    </Text>
+                  </Stack>
+                  <Stack horizontal style={{ marginLeft: "15px" }} tokens={{ childrenGap: 8 }}>
+                    <Text style={{ fontSize: "16px", lineHeight: "16px", verticalAlign: "middle" }}>•</Text>
+                    <Text style={{ verticalAlign: "middle" }}>Help answer questions about Cosmos DB</Text>
+                  </Stack>
+                </Stack>
+                <Text variant="medium">
+                  To get started, ask me a question or use one of the sample prompts to generate a query. AI-generated
+                  content may be incorrect.
+                </Text>
+              </Stack>
+              {chatMessages.map((message, index) => (
+                <Stack
+                  key={index}
+                  horizontalAlign="center"
+                  tokens={{ padding: 8, childrenGap: 8 }}
+                  style={{
+                    backgroundColor: "#E0E7FF",
+                    borderRadius: "8px",
+                    margin: "5px 10px",
+                    textAlign: "start",
+                  }}
+                >
+                  {message}
+                </Stack>
+              ))}
+            </div>
+
+            {chatMessages.length === 0 && (
+              <Stack
+                horizontalAlign="end"
+                verticalAlign="end"
+                style={{
+                  display: "flex",
+                  alignItems: "center",
+                  padding: "10px",
+                  margin: "10px",
+                }}
+              >
+                <Text
+                  onClick={() => handleInputChange(sampleChatMessages[0])}
+                  style={{
+                    cursor: "pointer",
+                    border: "1.5px solid #B0BEFF",
+                    width: "100%",
+                    padding: "2px",
+                    borderRadius: "4px",
+                    marginBottom: "5px",
+                  }}
+                >
+                  {sampleChatMessages[0]}
+                </Text>
+                <Text
+                  onClick={() => handleInputChange(sampleChatMessages[1])}
+                  style={{
+                    cursor: "pointer",
+                    border: "1.5px solid #B0BEFF",
+                    width: "100%",
+                    padding: "2px",
+                    borderRadius: "4px",
+                    marginBottom: "5px",
+                  }}
+                >
+                  {sampleChatMessages[1]}
+                </Text>
+              </Stack>
+            )}
+
+            <Stack
+              horizontal
+              horizontalAlign="end"
+              verticalAlign="end"
+              style={{
+                display: "flex",
+                alignItems: "center",
+                borderRadius: "20px",
+                background: "white",
+                padding: "5px",
+                margin: "5px",
+              }}
+            >
+              <Stack>
+                <Image src={HintIcon} styles={promptStyles} onClick={() => setIsSamplePromptsOpen(true)} />
+                <SamplePrompts sampleProps={sampleProps} />
+              </Stack>
+              <TextField
+                placeholder="Write your own prompt or ask a question"
+                value={userPrompt}
+                onChange={(_, newValue) => handleInputChange(newValue)}
+                onKeyDown={handleEnterKeyPress}
+                multiline
+                resizable={false}
+                styles={{
+                  root: {
+                    width: "100%",
+                    height: "80px",
+                    borderRadius: "20px",
+                    padding: "8px",
+                    border: "none",
+                    outline: "none",
+                    marginLeft: "10px",
+                  },
+                  fieldGroup: { border: "none" },
+                }}
+              />
+              <IconButton iconProps={{ iconName: "Send" }} onClick={handleSendMessage} />
+            </Stack>
+          </Stack>
+        </>
+      )}
+    </Stack>
+  );
+};

src/Explorer/QueryCopilot/Sidebar/WelcomeSidebarPopup.tsx

@@ -0,0 +1,127 @@
+import { Image, Link, PrimaryButton, Stack, Text } from "@fluentui/react";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
+import React from "react";
+import Database from "../../../../images/CopilotDatabase.svg";
+import Flash from "../../../../images/CopilotFlash.svg";
+import CopilotSidebarWelcomeIllustration from "../../../../images/CopilotSidebarWelcomeIllustration.svg";
+import Thumb from "../../../../images/CopilotThumb.svg";
+
+export const WelcomeSidebarPopup: React.FC = (): JSX.Element => {
+  const { setShowWelcomeSidebar } = useQueryCopilot();
+
+  const hideModal = () => {
+    setShowWelcomeSidebar(false);
+    window.localStorage.setItem("showWelcomeSidebar", "false");
+  };
+
+  React.useEffect(() => {
+    const showWelcomeSidebar = window.localStorage.getItem("showWelcomeSidebar");
+    setShowWelcomeSidebar(showWelcomeSidebar && showWelcomeSidebar === "false" ? false : true);
+  }, []);
+
+  return (
+    <Stack
+      style={{
+        width: "100%",
+        height: "100%",
+        overflow: "auto",
+        backgroundColor: "#FAFAFA",
+      }}
+    >
+      <div
+        style={{
+          margin: "20px 10px",
+          padding: "20px",
+          maxHeight: "100%",
+          boxSizing: "border-box",
+          borderRadius: "20px",
+          backgroundColor: "white",
+        }}
+      >
+        <Stack horizontalAlign="center" verticalAlign="center">
+          <Image src={CopilotSidebarWelcomeIllustration} />
+        </Stack>
+
+        <Stack>
+          <Stack.Item align="center" style={{ marginBottom: "10px" }}>
+            <Text className="title bold">Welcome to Copilot in CosmosDB</Text>
+          </Stack.Item>
+          <Stack.Item style={{ marginBottom: "15px" }}>
+            <Stack>
+              <Stack horizontal>
+                <Stack.Item align="start">
+                  <Image src={Flash} />
+                </Stack.Item>
+                <Stack.Item align="center" style={{ marginLeft: "10px" }}>
+                  <Text style={{ fontWeight: 600 }}>
+                    Let copilot do the work for you
+                    <br />
+                  </Text>
+                </Stack.Item>
+              </Stack>
+              <Stack.Item style={{ textAlign: "start", marginLeft: "25px" }}>
+                <Text>
+                  Ask Copilot to generate a query by describing the query in your words.
+                  <br />
+                  <Link href="http://aka.ms/cdb-copilot-learn-more">Learn more</Link>
+                </Text>
+              </Stack.Item>
+            </Stack>
+          </Stack.Item>
+          <Stack.Item style={{ marginBottom: "15px" }}>
+            <Stack>
+              <Stack horizontal>
+                <Stack.Item align="start">
+                  <Image src={Thumb} />
+                </Stack.Item>
+                <Stack.Item align="center" style={{ marginLeft: "10px" }}>
+                  <Text style={{ fontWeight: 600 }}>
+                    Use your judgement
+                    <br />
+                  </Text>
+                </Stack.Item>
+              </Stack>
+              <Stack.Item style={{ textAlign: "start", marginLeft: "25px" }}>
+                <Text>
+                  AI-generated content can have mistakes. Make sure it’s accurate and appropriate before using it.
+                  <br />
+                  <Link href="http://aka.ms/cdb-copilot-preview-terms">Read preview terms</Link>
+                </Text>
+              </Stack.Item>
+            </Stack>
+          </Stack.Item>
+          <Stack.Item style={{ marginBottom: "15px" }}>
+            <Stack>
+              <Stack horizontal>
+                <Stack.Item align="start">
+                  <Image src={Database} />
+                </Stack.Item>
+                <Stack.Item align="center" style={{ marginLeft: "10px" }}>
+                  <Text style={{ fontWeight: 600 }}>
+                    Copilot currently works only a sample database
+                    <br />
+                  </Text>
+                </Stack.Item>
+              </Stack>
+              <Stack.Item style={{ textAlign: "start", marginLeft: "25px" }}>
+                <Text>
+                  Copilot is setup on a sample database we have configured for you at no cost
+                  <br />
+                  <Link href="http://aka.ms/cdb-copilot-learn-more">Learn more</Link>
+                </Text>
+              </Stack.Item>
+            </Stack>
+          </Stack.Item>
+        </Stack>
+
+        <Stack>
+          <Stack.Item align="center">
+            <PrimaryButton style={{ width: "224px", height: "32px" }} onClick={hideModal}>
+              Get Started
+            </PrimaryButton>
+          </Stack.Item>
+        </Stack>
+      </div>
+    </Stack>
+  );
+};

src/Explorer/SplashScreen/SplashScreen.tsx

@@ -148,7 +148,13 @@ export class SplashScreen extends React.Component<SplashScreenProps> {
                 "Copilot is your AI buddy that helps you write Azure Cosmos DB queries like a pro. Try it using our sample data set now!"
               }
               onClick={() => {
-                useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot);
+                const copilotVersion = userContext.features.copilotVersion;
+                if (copilotVersion === "v1.0") {
+                  useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot);
+                } else if (copilotVersion === "v2.0") {
+                  const sampleCollection = useDatabases.getState().sampleDataResourceTokenCollection;
+                  sampleCollection.onNewQueryClick(sampleCollection, undefined);
+                }
                 traceOpen(Action.OpenQueryCopilotFromSplashScreen, { apiType: userContext.apiType });
               }}
             />

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -1,8 +1,12 @@
 import { FeedOptions } from "@azure/cosmos";
+import { QueryCopilotSidebar } from "Explorer/QueryCopilot/Sidebar/QueryCopilotSidebar";
 import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection";
+import { useDatabases } from "Explorer/useDatabases";
+import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
 import React, { Fragment } from "react";
 import SplitterLayout from "react-splitter-layout";
 import "react-splitter-layout/lib/index.css";
+import LaunchCopilot from "../../../../images/CopilotTabIcon.svg";
 import ExecuteQueryIcon from "../../../../images/ExecuteQuery.svg";
 import SaveQueryIcon from "../../../../images/save-cosmos.svg";
 import { NormalizedEventKey } from "../../../Common/Constants";
@@ -72,12 +76,15 @@ interface IQueryTabStates {
   error: string;
   isExecutionError: boolean;
   isExecuting: boolean;
+  showCopilotSidebar: boolean;
+  isCopilotTabActive: boolean;
 }
 
 export default class QueryTabComponent extends React.Component<IQueryTabComponentProps, IQueryTabStates> {
   public queryEditorId: string;
   public executeQueryButton: Button;
   public saveQueryButton: Button;
+  public launchCopilotButton: Button;
   public splitterId: string;
   public isPreferredApiMongoDB: boolean;
   public isCloseClicked: boolean;
@@ -94,6 +101,9 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       error: "",
       isExecutionError: this.props.isExecutionError,
       isExecuting: false,
+      showCopilotSidebar: useQueryCopilot.getState().showCopilotSidebar,
+      isCopilotTabActive:
+        useDatabases.getState().sampleDataResourceTokenCollection.databaseId === this.props.collection.databaseId,
     };
     this.isCloseClicked = false;
     this.splitterId = this.props.tabId + "_splitter";
@@ -111,6 +121,11 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       visible: isSaveQueryBtnEnabled,
     };
 
+    this.launchCopilotButton = {
+      enabled: userContext.apiType === "SQL" && true,
+      visible: userContext.apiType === "SQL" && true,
+    };
+
     this.props.tabsBaseInstance.updateNavbarWithTabsButtons();
     props.onTabAccessor({
       onTabClickEvent: this.onTabClick.bind(this),
@@ -121,6 +136,9 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
 
   public onCloseClick(isClicked: boolean): void {
     this.isCloseClicked = isClicked;
+    if (useQueryCopilot.getState().wasCopilotUsed && this.state.isCopilotTabActive) {
+      useQueryCopilot.getState().resetQueryCopilotStates();
+    }
   }
 
   public getCurrentEditorQuery(): string {
@@ -146,6 +164,10 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
     useSidePanel.getState().openSidePanel("Save Query", <SaveQueryPane explorer={this.props.collection.container} />);
   };
 
+  public launchQueryCopilotChat = (): void => {
+    useQueryCopilot.getState().setShowCopilotSidebar(!useQueryCopilot.getState().showCopilotSidebar);
+  };
+
   public onSavedQueriesClick = (): void => {
     useSidePanel
       .getState()
@@ -269,6 +291,18 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       });
     }
 
+    if (this.launchCopilotButton.visible && this.state.isCopilotTabActive) {
+      const label = "Launch Copilot";
+      buttons.push({
+        iconSrc: LaunchCopilot,
+        iconAlt: label,
+        onCommandClick: this.launchQueryCopilotChat,
+        commandButtonLabel: label,
+        ariaLabel: label,
+        hasPopup: false,
+      });
+    }
+
     return buttons;
   }
 
@@ -306,11 +340,23 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
     useCommandBar.getState().setContextButtons(this.getTabsButtons());
   }
 
+  private unsubscribeCopilotSidebar: () => void;
+
   componentDidMount(): void {
+    this.unsubscribeCopilotSidebar = useQueryCopilot.subscribe((state: QueryCopilotState) => {
+      if (this.state.showCopilotSidebar !== state.showCopilotSidebar) {
+        this.setState({ showCopilotSidebar: state.showCopilotSidebar });
+      }
+    });
+
     useCommandBar.getState().setContextButtons(this.getTabsButtons());
   }
 
-  render(): JSX.Element {
+  componentWillUnmount(): void {
+    this.unsubscribeCopilotSidebar();
+  }
+
+  private getEditorAndQueryResult(): JSX.Element {
     return (
       <Fragment>
         <div className="tab-pane" id={this.props.tabId} role="tabpanel">
@@ -343,4 +389,19 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       </Fragment>
     );
   }
+
+  render(): JSX.Element {
+    return (
+      <div style={{ display: "flex", flexDirection: "row", height: "100%" }}>
+        <div style={{ width: this.state.showCopilotSidebar ? "70%" : "100%", height: "100%" }}>
+          {this.getEditorAndQueryResult()}
+        </div>
+        {this.state.showCopilotSidebar && this.state.isCopilotTabActive && (
+          <div style={{ width: "30%", height: "100%" }}>
+            <QueryCopilotSidebar />
+          </div>
+        )}
+      </div>
+    );
+  }
 }

src/Platform/Hosted/extractFeatures.ts

@@ -37,6 +37,7 @@ export type Features = {
   readonly loadLegacyMongoShellFromBE: boolean;
   readonly enableCopilot: boolean;
   readonly enableNPSSurvey: boolean;
+  readonly copilotVersion?: string;
 
   // can be set via both flight and feature flag
   autoscaleDefault: boolean;
@@ -106,6 +107,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     loadLegacyMongoShellFromBE: "true" === get("loadlegacymongoshellfrombe"),
     enableCopilot: "true" === get("enablecopilot"),
     enableNPSSurvey: "true" === get("enablenpssurvey"),
+    copilotVersion: get("copilotVersion") ? get("copilotVersion") : "v1.0",
   };
 }
 

src/hooks/useQueryCopilot.ts

@@ -3,7 +3,7 @@ import { QueryResults } from "Contracts/ViewModels";
 import { guid } from "Explorer/Tables/Utilities";
 import create, { UseStore } from "zustand";
 
-interface QueryCopilotState {
+export interface QueryCopilotState {
   generatedQuery: string;
   likeQuery: boolean;
   userPrompt: string;
@@ -26,6 +26,10 @@ interface QueryCopilotState {
   showCopyPopup: boolean;
   showErrorMessageBar: boolean;
   generatedQueryComments: string;
+  wasCopilotUsed: boolean;
+  showWelcomeSidebar: boolean;
+  showCopilotSidebar: boolean;
+  chatMessages: string[];
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void;
   closeFeedbackModal: () => void;
@@ -50,6 +54,11 @@ interface QueryCopilotState {
   setshowCopyPopup: (showCopyPopup: boolean) => void;
   setShowErrorMessageBar: (showErrorMessageBar: boolean) => void;
   setGeneratedQueryComments: (generatedQueryComments: string) => void;
+  setWasCopilotUsed: (wasCopilotUsed: boolean) => void;
+  setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => void;
+  setShowCopilotSidebar: (showCopilotSidebar: boolean) => void;
+  setChatMessages: (chatMessages: string[]) => void;
+
   resetQueryCopilotStates: () => void;
 }
 
@@ -78,6 +87,10 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   showCopyPopup: false,
   showErrorMessageBar: false,
   generatedQueryComments: "",
+  wasCopilotUsed: false,
+  showWelcomeSidebar: true,
+  showCopilotSidebar: false,
+  chatMessages: [],
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
     set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
@@ -104,6 +117,10 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   setshowCopyPopup: (showCopyPopup: boolean) => set({ showCopyPopup }),
   setShowErrorMessageBar: (showErrorMessageBar: boolean) => set({ showErrorMessageBar }),
   setGeneratedQueryComments: (generatedQueryComments: string) => set({ generatedQueryComments }),
+  setWasCopilotUsed: (wasCopilotUsed: boolean) => set({ wasCopilotUsed }),
+  setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => set({ showWelcomeSidebar }),
+  setShowCopilotSidebar: (showCopilotSidebar: boolean) => set({ showCopilotSidebar }),
+  setChatMessages: (chatMessages: string[]) => set({ chatMessages }),
 
   resetQueryCopilotStates: () => {
     set((state) => ({
@@ -130,6 +147,9 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       showCopyPopup: false,
       showErrorMessageBar: false,
       generatedQueryComments: "",
+      wasCopilotUsed: false,
+      showCopilotSidebar: false,
+      chatMessages: [],
     }));
   },
 }));