From 0e124f4881eaaaddec04dfb2cee48aa7c453cfa5 Mon Sep 17 00:00:00 2001 From: sunghyunkang1111 <114709653+sunghyunkang1111@users.noreply.github.com> Date: Thu, 9 Nov 2023 11:55:25 -0600 Subject: [PATCH] Copilot user db (#1672) * Implement copilot for user database * Fix minor bugs * fix bugs * Add user database copilot * Add placeholder text on copilot * Add AFEC adn killswitch * Add new v2 sampledatabase endpoint * Add telemetry * fix telemetry bug * Add query edited telemetry * add authorization header * Add back to the staging env for phoenix * point to stage for phoenix * Preview commit for test env * Preview link for staging * change the staging url * fix lint, unit tests * fix lint, unit tests * fix formatting --- src/Common/Constants.ts | 1 + src/ConfigContext.ts | 3 +- src/Contracts/DataModels.ts | 16 +- src/Explorer/Controls/Editor/EditorReact.tsx | 4 +- src/Explorer/Explorer.tsx | 72 +++++--- .../CommandBarComponentButtonFactory.tsx | 12 +- .../Menus/CommandBar/CommandBarUtil.tsx | 7 +- .../Modal/QueryCopilotFeedbackModal.test.tsx | 120 +++++++++++-- .../Modal/QueryCopilotFeedbackModal.tsx | 21 ++- .../QueryCopilot/Modal/WelcomeModal.tsx | 37 +--- .../QueryCopilotFeedbackModal.test.tsx.snap | 15 -- .../__snapshots__/WelcomeModal.test.tsx.snap | 51 +----- .../QueryCopilot/QueryCopilotContext.tsx | 135 ++++++++++++++ .../QueryCopilot/QueryCopilotPromptbar.tsx | 81 +++++---- src/Explorer/QueryCopilot/QueryCopilotTab.tsx | 25 ++- .../QueryCopilot/QueryCopilotUtilities.ts | 21 +++ .../Shared/QueryCopilotClient.test.ts | 14 +- .../QueryCopilot/Shared/QueryCopilotClient.ts | 169 +++++++++++++++++- .../Shared/QueryCopilotInterfaces.ts | 5 + .../Shared/QueryCopilotResults.tsx | 2 +- .../QueryCopilotTab.test.tsx.snap | 2 + src/Explorer/SplashScreen/SplashScreen.tsx | 13 +- src/Explorer/Tabs/DocumentsTab.ts | 5 +- src/Explorer/Tabs/QueryTab/QueryTab.tsx | 17 +- .../Tabs/QueryTab/QueryTabComponent.test.tsx | 50 +++++- .../Tabs/QueryTab/QueryTabComponent.tsx | 158 ++++++++++++++-- src/Explorer/Tabs/Tabs.tsx | 4 +- src/Explorer/Tabs/TabsBase.ts | 6 +- src/Explorer/Tree/Collection.ts | 4 +- src/Explorer/Tree/ResourceTokenCollection.ts | 2 +- src/Explorer/Tree/ResourceTree.tsx | 4 +- src/Explorer/useDatabases.ts | 8 +- src/Main.tsx | 22 +-- src/Platform/Hosted/extractFeatures.ts | 2 +- src/Shared/Telemetry/TelemetryConstants.ts | 4 + src/hooks/useKnockoutExplorer.ts | 18 +- src/hooks/useQueryCopilot.ts | 31 +++- src/hooks/useTabs.ts | 2 +- 38 files changed, 911 insertions(+), 252 deletions(-) create mode 100644 src/Explorer/QueryCopilot/QueryCopilotContext.tsx diff --git a/src/Common/Constants.ts b/src/Common/Constants.ts index 46b890d29..1160c2296 100644 --- a/src/Common/Constants.ts +++ b/src/Common/Constants.ts @@ -171,6 +171,7 @@ export class Areas { public static Tab: string = "Tab"; public static ShareDialog: string = "Share Access Dialog"; public static Notebook: string = "Notebook"; + public static Copilot: string = "Copilot"; } export class HttpHeaders { diff --git a/src/ConfigContext.ts b/src/ConfigContext.ts index 2c11d5daa..517b3f200 100644 --- a/src/ConfigContext.ts +++ b/src/ConfigContext.ts @@ -1,3 +1,4 @@ +import { JunoEndpoints } from "Common/Constants"; import { allowedAadEndpoints, allowedArcadiaEndpoints, @@ -78,7 +79,7 @@ let configContext: Readonly = { ARCADIA_LIVY_ENDPOINT_DNS_ZONE: "dev.azuresynapse.net", GITHUB_CLIENT_ID: "6cb2f63cf6f7b5cbdeca", // Registered OAuth app: https://github.com/organizations/AzureCosmosDBNotebooks/settings/applications/1189306 GITHUB_TEST_ENV_CLIENT_ID: "b63fc8cbf87fd3c6e2eb", // Registered OAuth app: https://github.com/organizations/AzureCosmosDBNotebooks/settings/applications/1777772 - JUNO_ENDPOINT: "https://tools.cosmos.azure.com", + JUNO_ENDPOINT: JunoEndpoints.Prod, BACKEND_ENDPOINT: "https://main.documentdb.ext.azure.com", isTerminalEnabled: false, isPhoenixEnabled: false, diff --git a/src/Contracts/DataModels.ts b/src/Contracts/DataModels.ts index c2ec21519..ac7b9499e 100644 --- a/src/Contracts/DataModels.ts +++ b/src/Contracts/DataModels.ts @@ -457,8 +457,11 @@ export interface ContainerInfo { } export interface IProvisionData { - cosmosEndpoint: string; + cosmosEndpoint?: string; poolId: string; + databaseId?: string; + containerId?: string; + mode?: string; } export interface IContainerData { @@ -601,3 +604,14 @@ export enum PhoenixErrorType { PhoenixFlightFallback = "PhoenixFlightFallback", UserMissingPermissionsError = "UserMissingPermissionsError", } + +export interface CopilotEnabledConfiguration { + isEnabled: boolean; +} + +export interface FeatureRegistration { + name: string; + properties: { + state: string; + }; +} diff --git a/src/Explorer/Controls/Editor/EditorReact.tsx b/src/Explorer/Controls/Editor/EditorReact.tsx index 7c8917a56..7e3ff732b 100644 --- a/src/Explorer/Controls/Editor/EditorReact.tsx +++ b/src/Explorer/Controls/Editor/EditorReact.tsx @@ -48,7 +48,7 @@ export class EditorReact extends React.Component { + public async allocateContainer(poolId: PoolIdType, mode?: string): Promise { const shouldUseNotebookStates = poolId === PoolIdType.DefaultPoolId ? true : false; const notebookServerInfo = shouldUseNotebookStates ? useNotebook.getState().notebookServerInfo @@ -425,10 +426,6 @@ export default class Explorer { (notebookServerInfo === undefined || (notebookServerInfo && notebookServerInfo.notebookServerEndpoint === undefined)) ) { - const provisionData: IProvisionData = { - cosmosEndpoint: userContext?.databaseAccount?.properties?.documentEndpoint, - poolId: shouldUseNotebookStates ? undefined : poolId, - }; const connectionStatus: ContainerConnectionInfo = { status: ConnectionStatusType.Connecting, }; @@ -436,14 +433,26 @@ export default class Explorer { shouldUseNotebookStates && useNotebook.getState().setConnectionInfo(connectionStatus); let connectionInfo; + let provisionData: IProvisionData; try { TelemetryProcessor.traceStart(Action.PhoenixConnection, { dataExplorerArea: Areas.Notebook, }); - shouldUseNotebookStates - ? useNotebook.getState().setIsAllocating(true) - : useQueryCopilot.getState().setIsAllocatingContainer(true); - + if (shouldUseNotebookStates) { + useNotebook.getState().setIsAllocating(true); + provisionData = { + cosmosEndpoint: userContext?.databaseAccount?.properties?.documentEndpoint, + poolId: undefined, + }; + } else { + useQueryCopilot.getState().setIsAllocatingContainer(true); + provisionData = { + poolId: poolId, + databaseId: useTabs.getState().activeTab.collection.databaseId, + containerId: useTabs.getState().activeTab.collection.id(), + mode: mode, + }; + } connectionInfo = await this.phoenixClient.allocateContainer(provisionData); if (!connectionInfo?.data?.phoenixServiceUrl) { throw new Error(`PhoenixServiceUrl is invalid!`); @@ -459,19 +468,21 @@ export default class Explorer { error: getErrorMessage(error), errorStack: getErrorStack(error), }); - connectionStatus.status = ConnectionStatusType.Failed; - shouldUseNotebookStates - ? useNotebook.getState().resetContainerConnection(connectionStatus) - : useQueryCopilot.getState().resetContainerConnection(); - if (error?.status === HttpStatusCodes.Forbidden && error.message) { - useDialog.getState().showOkModalDialog("Connection Failed", `${error.message}`); - } else { - useDialog - .getState() - .showOkModalDialog( - "Connection Failed", - "We are unable to connect to the temporary workspace. Please try again in a few minutes. If the error persists, file a support ticket.", - ); + if (shouldUseNotebookStates) { + connectionStatus.status = ConnectionStatusType.Failed; + shouldUseNotebookStates + ? useNotebook.getState().resetContainerConnection(connectionStatus) + : useQueryCopilot.getState().resetContainerConnection(); + if (error?.status === HttpStatusCodes.Forbidden && error.message) { + useDialog.getState().showOkModalDialog("Connection Failed", `${error.message}`); + } else { + useDialog + .getState() + .showOkModalDialog( + "Connection Failed", + "We are unable to connect to the temporary workspace. Please try again in a few minutes. If the error persists, file a support ticket.", + ); + } } throw error; } finally { @@ -485,11 +496,11 @@ export default class Explorer { } } - private async setNotebookInfo( + public async setNotebookInfo( shouldUseNotebookStates: boolean, connectionInfo: IResponse, connectionStatus: DataModels.ContainerConnectionInfo, - ) { + ): Promise { const containerData = { forwardingId: connectionInfo.data.forwardingId, dbAccountName: userContext.databaseAccount.name, @@ -510,6 +521,7 @@ export default class Explorer { shouldUseNotebookStates ? useNotebook.getState().setNotebookServerInfo(noteBookServerInfo) : useQueryCopilot.getState().setNotebookServerInfo(noteBookServerInfo); + shouldUseNotebookStates && this.notebookManager?.notebookClient .getMemoryUsage() @@ -1372,6 +1384,16 @@ export default class Explorer { await this.refreshSampleData(); } + public async configureCopilot(): Promise { + if (userContext.apiType !== "SQL") { + return; + } + const copilotEnabled = await getCopilotEnabled(); + const copilotUserDBEnabled = await isCopilotFeatureRegistered(userContext.subscriptionId); + useQueryCopilot.getState().setCopilotEnabled(copilotEnabled); + useQueryCopilot.getState().setCopilotUserDBEnabled(copilotUserDBEnabled); + } + public async refreshSampleData(): Promise { if (!userContext.sampleDataConnectionInfo) { return; diff --git a/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx b/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx index 94c1726f6..5e9de9cde 100644 --- a/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx +++ b/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx @@ -1,6 +1,3 @@ -import { Action } from "Shared/Telemetry/TelemetryConstants"; -import { traceOpen } from "Shared/Telemetry/TelemetryProcessor"; -import { ReactTabKind, useTabs } from "hooks/useTabs"; import * as React from "react"; import AddCollectionIcon from "../../../../images/AddCollection.svg"; import AddDatabaseIcon from "../../../../images/AddDatabase.svg"; @@ -337,13 +334,8 @@ function createNewSQLQueryButton(selectedNodeState: SelectedNodeState): CommandB iconSrc: AddSqlQueryIcon, iconAlt: label, onCommandClick: () => { - if (useSelectedNode.getState().isQueryCopilotCollectionSelected()) { - useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot); - traceOpen(Action.OpenQueryCopilotFromNewQuery, { apiType: userContext.apiType }); - } else { - const selectedCollection: ViewModels.Collection = selectedNodeState.findSelectedCollection(); - selectedCollection && selectedCollection.onNewQueryClick(selectedCollection); - } + const selectedCollection: ViewModels.Collection = selectedNodeState.findSelectedCollection(); + selectedCollection && selectedCollection.onNewQueryClick(selectedCollection); }, commandButtonLabel: label, ariaLabel: label, diff --git a/src/Explorer/Menus/CommandBar/CommandBarUtil.tsx b/src/Explorer/Menus/CommandBar/CommandBarUtil.tsx index c6bef80ac..570a7ad1f 100644 --- a/src/Explorer/Menus/CommandBar/CommandBarUtil.tsx +++ b/src/Explorer/Menus/CommandBar/CommandBarUtil.tsx @@ -6,6 +6,7 @@ import { IDropdownOption, IDropdownStyles, } from "@fluentui/react"; +import { useQueryCopilot } from "hooks/useQueryCopilot"; import * as React from "react"; import _ from "underscore"; import ChevronDownIcon from "../../../../images/Chevron_down.svg"; @@ -57,7 +58,11 @@ export const convertButton = (btns: CommandButtonComponentProps[], backgroundCol }, onClick: (ev?: React.MouseEvent | React.KeyboardEvent) => { btn.onCommandClick(ev); - TelemetryProcessor.trace(Action.ClickCommandBarButton, ActionModifiers.Mark, { label }); + let copilotEnabled = false; + if (useQueryCopilot.getState().copilotEnabled && useQueryCopilot.getState().copilotUserDBEnabled) { + copilotEnabled = useQueryCopilot.getState().copilotEnabledforExecution; + } + TelemetryProcessor.trace(Action.ClickCommandBarButton, ActionModifiers.Mark, { label, copilotEnabled }); }, key: `${btn.commandButtonLabel}${index}`, text: label, diff --git a/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx b/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx index 14cca91a7..4c805ad36 100644 --- a/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx +++ b/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx @@ -1,10 +1,10 @@ import { Checkbox, ChoiceGroup, DefaultButton, IconButton, PrimaryButton, TextField } from "@fluentui/react"; import Explorer from "Explorer/Explorer"; import { QueryCopilotFeedbackModal } from "Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal"; +import { useCopilotStore } from "Explorer/QueryCopilot/QueryCopilotContext"; import { SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient"; import { getUserEmail } from "Utils/UserUtils"; import { shallow } from "enzyme"; -import { useQueryCopilot } from "hooks/useQueryCopilot"; import React from "react"; jest.mock("Utils/UserUtils"); @@ -13,21 +13,49 @@ jest.mock("Utils/UserUtils"); jest.mock("Explorer/QueryCopilot/Shared/QueryCopilotClient"); SubmitFeedback as jest.Mock; +jest.mock("Explorer/QueryCopilot/QueryCopilotContext"); +const mockUseCopilotStore = useCopilotStore as jest.Mock; +const mockReturnValue = { + generatedQuery: "test query", + userPrompt: "test prompt", + likeQuery: false, + showFeedbackModal: false, + closeFeedbackModal: jest.fn, + setHideFeedbackModalForLikedQueries: jest.fn, +}; + describe("Query Copilot Feedback Modal snapshot test", () => { beforeEach(() => { + mockUseCopilotStore.mockReturnValue(mockReturnValue); jest.clearAllMocks(); }); it("shoud render and match snapshot", () => { - useQueryCopilot.getState().openFeedbackModal("test query", false, "test prompt"); - - const wrapper = shallow(); + mockUseCopilotStore.mockReturnValue({ + ...mockReturnValue, + showFeedbackModal: true, + }); + const wrapper = shallow( + , + ); expect(wrapper.props().isOpen).toBeTruthy(); expect(wrapper).toMatchSnapshot(); }); it("should close on cancel click", () => { - const wrapper = shallow(); + const wrapper = shallow( + , + ); const cancelButton = wrapper.find(IconButton); cancelButton.simulate("click"); @@ -38,7 +66,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => { }); it("should get user unput", () => { - const wrapper = shallow(); + const wrapper = shallow( + , + ); const testUserInput = "test user input"; const userInput = wrapper.find(TextField).first(); @@ -49,7 +84,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => { }); it("should record user contact choice no", () => { - const wrapper = shallow(); + const wrapper = shallow( + , + ); const contactAllowed = wrapper.find(ChoiceGroup); contactAllowed.simulate("change", {}, { key: "no" }); @@ -60,7 +102,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => { }); it("should record user contact choice yes", () => { - const wrapper = shallow(); + const wrapper = shallow( + , + ); const contactAllowed = wrapper.find(ChoiceGroup); contactAllowed.simulate("change", {}, { key: "yes" }); @@ -71,7 +120,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => { }); it("should not render dont show again button", () => { - const wrapper = shallow(); + const wrapper = shallow( + , + ); const dontShowAgain = wrapper.find(Checkbox); @@ -80,8 +136,19 @@ describe("Query Copilot Feedback Modal snapshot test", () => { }); it("should render dont show again button and check it ", () => { - useQueryCopilot.getState().openFeedbackModal("test query", true, "test prompt"); - const wrapper = shallow(); + mockUseCopilotStore.mockReturnValue({ + ...mockReturnValue, + showFeedbackModal: true, + likeQuery: true, + }); + const wrapper = shallow( + , + ); const dontShowAgain = wrapper.find(Checkbox); dontShowAgain.simulate("change", {}, true); @@ -92,7 +159,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => { }); it("should cancel submission", () => { - const wrapper = shallow(); + const wrapper = shallow( + , + ); const cancelButton = wrapper.find(DefaultButton); cancelButton.simulate("click"); @@ -104,7 +178,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => { it("should not submit submission if required description field is null", () => { const explorer = new Explorer(); - const wrapper = shallow(); + const wrapper = shallow( + , + ); const submitButton = wrapper.find(PrimaryButton); submitButton.simulate("click"); @@ -114,9 +195,15 @@ describe("Query Copilot Feedback Modal snapshot test", () => { }); it("should submit submission", () => { - useQueryCopilot.getState().openFeedbackModal("test query", false, "test prompt"); const explorer = new Explorer(); - const wrapper = shallow(); + const wrapper = shallow( + , + ); const submitButton = wrapper.find("form"); submitButton.simulate("submit"); @@ -124,6 +211,9 @@ describe("Query Copilot Feedback Modal snapshot test", () => { expect(SubmitFeedback).toHaveBeenCalledTimes(1); expect(SubmitFeedback).toHaveBeenCalledWith({ + containerId: "CopilotUserContainer", + databaseId: "CopilotUserDb", + mode: "User", params: { likeQuery: false, generatedQuery: "test query", diff --git a/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx b/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx index aa865ed0e..6cea9ff68 100644 --- a/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx +++ b/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx @@ -11,12 +11,22 @@ import { TextField, } from "@fluentui/react"; import Explorer from "Explorer/Explorer"; +import { useCopilotStore } from "Explorer/QueryCopilot/QueryCopilotContext"; import { SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient"; -import { useQueryCopilot } from "hooks/useQueryCopilot"; import React from "react"; import { getUserEmail } from "../../../Utils/UserUtils"; -export const QueryCopilotFeedbackModal = ({ explorer }: { explorer: Explorer }): JSX.Element => { +export const QueryCopilotFeedbackModal = ({ + explorer, + databaseId, + containerId, + mode, +}: { + explorer: Explorer; + databaseId: string; + containerId: string; + mode: string; +}): JSX.Element => { const { generatedQuery, userPrompt, @@ -24,7 +34,7 @@ export const QueryCopilotFeedbackModal = ({ explorer }: { explorer: Explorer }): showFeedbackModal, closeFeedbackModal, setHideFeedbackModalForLikedQueries, - } = useQueryCopilot(); + } = useCopilotStore(); const [isContactAllowed, setIsContactAllowed] = React.useState(false); const [description, setDescription] = React.useState(""); const [doNotShowAgainChecked, setDoNotShowAgainChecked] = React.useState(false); @@ -35,7 +45,10 @@ export const QueryCopilotFeedbackModal = ({ explorer }: { explorer: Explorer }): setHideFeedbackModalForLikedQueries(doNotShowAgainChecked); SubmitFeedback({ params: { generatedQuery, likeQuery, description, userPrompt, contact }, - explorer: explorer, + explorer, + databaseId, + containerId, + mode: mode, }); }; diff --git a/src/Explorer/QueryCopilot/Modal/WelcomeModal.tsx b/src/Explorer/QueryCopilot/Modal/WelcomeModal.tsx index da45ff904..bec25b488 100644 --- a/src/Explorer/QueryCopilot/Modal/WelcomeModal.tsx +++ b/src/Explorer/QueryCopilot/Modal/WelcomeModal.tsx @@ -1,7 +1,6 @@ import { IconButton, Image, Link, Modal, PrimaryButton, Stack, StackItem, Text } from "@fluentui/react"; import { useBoolean } from "@fluentui/react-hooks"; import React from "react"; -import Database from "../../../../images/CopilotDatabase.svg"; import Flash from "../../../../images/CopilotFlash.svg"; import Thumb from "../../../../images/CopilotThumb.svg"; import CoplilotWelcomeIllustration from "../../../../images/CopliotWelcomeIllustration.svg"; @@ -23,8 +22,10 @@ export const WelcomeModal = ({ visible }: { visible: boolean }): JSX.Element => onDismiss={hideModal} isBlocking={false} styles={{ - scrollableContent: { - minHeight: 680, + main: { + maxHeight: 530, + borderRadius: 10, + overflow: "hidden", }, }} > @@ -52,7 +53,7 @@ export const WelcomeModal = ({ visible }: { visible: boolean }): JSX.Element => - Welcome to Copilot in Azure Cosmos DB (Private Preview) + Welcome to Copilot in Azure Cosmos DB @@ -69,7 +70,7 @@ export const WelcomeModal = ({ visible }: { visible: boolean }): JSX.Element => Ask Copilot to generate a query by describing the query in your words.
- + Learn more @@ -87,31 +88,11 @@ export const WelcomeModal = ({ visible }: { visible: boolean }): JSX.Element => - AI-generated content can have mistakes. Make sure it’s accurate and appropriate before using it. + AI-generated content can have mistakes. Make sure it is accurate and appropriate before executing the + query.
- Read preview terms - - - - - - - - - - - Query Copilot works on a sample database. -
- - - - - While in Private Preview, Query Copilot is setup to work on sample database we have configured for you - at no cost. -
- - Learn more + Read our preview terms here diff --git a/src/Explorer/QueryCopilot/Modal/__snapshots__/QueryCopilotFeedbackModal.test.tsx.snap b/src/Explorer/QueryCopilot/Modal/__snapshots__/QueryCopilotFeedbackModal.test.tsx.snap index 2d333adb5..d0b2a3481 100644 --- a/src/Explorer/QueryCopilot/Modal/__snapshots__/QueryCopilotFeedbackModal.test.tsx.snap +++ b/src/Explorer/QueryCopilot/Modal/__snapshots__/QueryCopilotFeedbackModal.test.tsx.snap @@ -291,21 +291,6 @@ exports[`Query Copilot Feedback Modal snapshot test should cancel submission 1`] for more information. - - Welcome to Copilot in Azure Cosmos DB (Private Preview) + Welcome to Copilot in Azure Cosmos DB Learn more @@ -143,50 +145,13 @@ exports[`Query Copilot Welcome Modal snapshot test should render when isOpen is - AI-generated content can have mistakes. Make sure it’s accurate and appropriate before using it. + AI-generated content can have mistakes. Make sure it is accurate and appropriate before executing the query.
- Read preview terms - - - - - - - - - - - Query Copilot works on a sample database. -
- - - - - While in Private Preview, Query Copilot is setup to work on sample database we have configured for you at no cost. -
- - Learn more + Read our preview terms here diff --git a/src/Explorer/QueryCopilot/QueryCopilotContext.tsx b/src/Explorer/QueryCopilot/QueryCopilotContext.tsx new file mode 100644 index 000000000..1ec2530d5 --- /dev/null +++ b/src/Explorer/QueryCopilot/QueryCopilotContext.tsx @@ -0,0 +1,135 @@ +import { MinimalQueryIterator } from "Common/IteratorUtilities"; +import { QueryResults } from "Contracts/ViewModels"; +import { CopilotMessage } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces"; +import { guid } from "Explorer/Tables/Utilities"; +import { QueryCopilotState } from "hooks/useQueryCopilot"; + +import React, { createContext, useContext, useState } from "react"; +import create from "zustand"; +const context = createContext(null); +const useCopilotStore = (): Partial => useContext(context); + +const CopilotProvider = ({ children }: { children: React.ReactNode }): JSX.Element => { + const [useStore] = useState(() => + create((set, get) => ({ + generatedQuery: "", + likeQuery: false, + userPrompt: "", + showFeedbackModal: false, + hideFeedbackModalForLikedQueries: false, + correlationId: "", + query: "SELECT * FROM c", + selectedQuery: "", + isGeneratingQuery: false, + isGeneratingExplanation: false, + isExecuting: false, + dislikeQuery: undefined, + showCallout: false, + showSamplePrompts: false, + queryIterator: undefined, + queryResults: undefined, + errorMessage: "", + isSamplePromptsOpen: false, + showDeletePopup: false, + showFeedbackBar: false, + showCopyPopup: false, + showErrorMessageBar: false, + showInvalidQueryMessageBar: false, + generatedQueryComments: "", + wasCopilotUsed: false, + showWelcomeSidebar: true, + showCopilotSidebar: false, + chatMessages: [], + shouldIncludeInMessages: true, + showExplanationBubble: false, + isAllocatingContainer: false, + + openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => + set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }), + closeFeedbackModal: () => set({ showFeedbackModal: false }), + setHideFeedbackModalForLikedQueries: (hideFeedbackModalForLikedQueries: boolean) => + set({ hideFeedbackModalForLikedQueries }), + refreshCorrelationId: () => set({ correlationId: guid() }), + setUserPrompt: (userPrompt: string) => set({ userPrompt }), + setQuery: (query: string) => set({ query }), + 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 }), + setShowCallout: (showCallout: boolean) => set({ showCallout }), + setShowSamplePrompts: (showSamplePrompts: boolean) => set({ showSamplePrompts }), + setQueryIterator: (queryIterator: MinimalQueryIterator | undefined) => set({ queryIterator }), + setQueryResults: (queryResults: QueryResults | undefined) => set({ queryResults }), + setErrorMessage: (errorMessage: string) => set({ errorMessage }), + setIsSamplePromptsOpen: (isSamplePromptsOpen: boolean) => set({ isSamplePromptsOpen }), + setShowDeletePopup: (showDeletePopup: boolean) => set({ showDeletePopup }), + setShowFeedbackBar: (showFeedbackBar: boolean) => set({ showFeedbackBar }), + setshowCopyPopup: (showCopyPopup: boolean) => set({ showCopyPopup }), + setShowErrorMessageBar: (showErrorMessageBar: boolean) => set({ showErrorMessageBar }), + setShowInvalidQueryMessageBar: (showInvalidQueryMessageBar: boolean) => set({ showInvalidQueryMessageBar }), + setGeneratedQueryComments: (generatedQueryComments: string) => set({ generatedQueryComments }), + setWasCopilotUsed: (wasCopilotUsed: boolean) => set({ wasCopilotUsed }), + setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => set({ showWelcomeSidebar }), + setShowCopilotSidebar: (showCopilotSidebar: boolean) => set({ showCopilotSidebar }), + setChatMessages: (chatMessages: CopilotMessage[]) => set({ chatMessages }), + setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => set({ shouldIncludeInMessages }), + setShowExplanationBubble: (showExplanationBubble: boolean) => set({ showExplanationBubble }), + + getState: () => { + return get(); + }, + + resetQueryCopilotStates: () => { + set((state) => ({ + ...state, + generatedQuery: "", + likeQuery: false, + userPrompt: "", + showFeedbackModal: false, + hideFeedbackModalForLikedQueries: false, + correlationId: "", + query: "SELECT * FROM c", + selectedQuery: "", + isGeneratingQuery: false, + isGeneratingExplanation: false, + isExecuting: false, + dislikeQuery: undefined, + showCallout: false, + showSamplePrompts: false, + queryIterator: undefined, + queryResults: undefined, + errorMessage: "", + isSamplePromptsOpen: false, + showDeletePopup: false, + showFeedbackBar: false, + showCopyPopup: false, + showErrorMessageBar: false, + showInvalidQueryMessageBar: false, + generatedQueryComments: "", + wasCopilotUsed: false, + showCopilotSidebar: false, + chatMessages: [], + shouldIncludeInMessages: true, + showExplanationBubble: false, + notebookServerInfo: { + notebookServerEndpoint: undefined, + authToken: undefined, + forwardingId: undefined, + }, + containerStatus: { + status: undefined, + durationLeftInMinutes: undefined, + phoenixServerInfo: undefined, + }, + isAllocatingContainer: false, + })); + }, + })), + ); + return {children}; +}; + +export { CopilotProvider, useCopilotStore }; diff --git a/src/Explorer/QueryCopilot/QueryCopilotPromptbar.tsx b/src/Explorer/QueryCopilot/QueryCopilotPromptbar.tsx index 88f9352ab..0a2752599 100644 --- a/src/Explorer/QueryCopilot/QueryCopilotPromptbar.tsx +++ b/src/Explorer/QueryCopilot/QueryCopilotPromptbar.tsx @@ -1,3 +1,5 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +/* eslint-disable no-console */ import { Callout, CommandBarButton, @@ -17,20 +19,20 @@ import { TextField, } from "@fluentui/react"; import { useBoolean } from "@fluentui/react-hooks"; -import { - ContainerStatusType, - PoolIdType, - QueryCopilotSampleContainerSchema, - ShortenedQueryCopilotSampleContainerSchema, -} from "Common/Constants"; import { handleError } from "Common/ErrorHandlingUtils"; import { createUri } from "Common/UrlUtility"; import { WelcomeModal } from "Explorer/QueryCopilot/Modal/WelcomeModal"; import { CopyPopup } from "Explorer/QueryCopilot/Popup/CopyPopup"; import { DeletePopup } from "Explorer/QueryCopilot/Popup/DeletePopup"; -import { SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient"; +import { + SuggestedPrompt, + getSampleDatabaseSuggestedPrompts, + getSuggestedPrompts, +} from "Explorer/QueryCopilot/QueryCopilotUtilities"; +import { SubmitFeedback, allocatePhoenixContainer } from "Explorer/QueryCopilot/Shared/QueryCopilotClient"; import { GenerateSQLQueryResponse, QueryCopilotProps } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces"; import { SamplePrompts, SamplePromptsProps } from "Explorer/QueryCopilot/Shared/SamplePrompts/SamplePrompts"; +import { Action } from "Shared/Telemetry/TelemetryConstants"; import { userContext } from "UserContext"; import { useQueryCopilot } from "hooks/useQueryCopilot"; import React, { useRef, useState } from "react"; @@ -38,17 +40,17 @@ import HintIcon from "../../../images/Hint.svg"; import CopilotIcon from "../../../images/QueryCopilotNewLogo.svg"; import RecentIcon from "../../../images/Recent.svg"; import errorIcon from "../../../images/close-black.svg"; +import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor"; import { useTabs } from "../../hooks/useTabs"; +import { useCopilotStore } from "../QueryCopilot/QueryCopilotContext"; +import { useSelectedNode } from "../useSelectedNode"; type QueryCopilotPromptProps = QueryCopilotProps & { + databaseId: string; + containerId: string; toggleCopilot: (toggle: boolean) => void; }; -interface SuggestedPrompt { - id: number; - text: string; -} - const promptStyles: IButtonStyles = { root: { border: 0, selectors: { ":hover": { outline: "1px dashed #605e5c" } } }, label: { fontWeight: 400, textAlign: "left", paddingLeft: 8 }, @@ -57,10 +59,13 @@ const promptStyles: IButtonStyles = { export const QueryCopilotPromptbar: React.FC = ({ explorer, toggleCopilot, + databaseId, + containerId, }: QueryCopilotPromptProps): JSX.Element => { const [copilotTeachingBubbleVisible, { toggle: toggleCopilotTeachingBubbleVisible }] = useBoolean(false); const inputEdited = useRef(false); const { + openFeedbackModal, hideFeedbackModalForLikedQueries, userPrompt, setUserPrompt, @@ -93,7 +98,7 @@ export const QueryCopilotPromptbar: React.FC = ({ setGeneratedQueryComments, setQueryResults, setErrorMessage, - } = useQueryCopilot(); + } = useCopilotStore(); const sampleProps: SamplePromptsProps = { isSamplePromptsOpen: isSamplePromptsOpen, @@ -118,14 +123,13 @@ export const QueryCopilotPromptbar: React.FC = ({ }, 6000); }; + const isSampleCopilotActive = useSelectedNode.getState().isQueryCopilotCollectionSelected(); const cachedHistoriesString = localStorage.getItem(`${userContext.databaseAccount?.id}-queryCopilotHistories`); const cachedHistories = cachedHistoriesString?.split("|"); const [histories, setHistories] = useState(cachedHistories || []); - const suggestedPrompts: SuggestedPrompt[] = [ - { id: 1, text: 'Show all products that have the word "ultra" in the name or description' }, - { id: 2, text: "What are all of the possible categories for the products, and their counts?" }, - { id: 3, text: 'Show me all products that have been reviewed by someone with a username that contains "bob"' }, - ]; + const suggestedPrompts: SuggestedPrompt[] = isSampleCopilotActive + ? getSampleDatabaseSuggestedPrompts() + : getSuggestedPrompts(); const [filteredHistories, setFilteredHistories] = useState(histories); const [filteredSuggestedPrompts, setFilteredSuggestedPrompts] = useState(suggestedPrompts); @@ -176,16 +180,11 @@ export const QueryCopilotPromptbar: React.FC = ({ setShowDeletePopup(false); useTabs.getState().setIsTabExecuting(true); useTabs.getState().setIsQueryErrorThrown(false); - if ( - useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active && - !userContext.features.disableCopilotPhoenixGateaway - ) { - await explorer.allocateContainer(PoolIdType.QueryCopilot); - } + const mode: string = isSampleCopilotActive ? "Sample" : "User"; + + await allocatePhoenixContainer({ explorer, databaseId, containerId, mode }); + const payload = { - containerSchema: userContext.features.enableCopilotFullSchema - ? QueryCopilotSampleContainerSchema - : ShortenedQueryCopilotSampleContainerSchema, userPrompt: userPrompt, }; useQueryCopilot.getState().refreshCorrelationId(); @@ -198,6 +197,7 @@ export const QueryCopilotPromptbar: React.FC = ({ headers: { "content-type": "application/json", "x-ms-correlationid": useQueryCopilot.getState().correlationId, + Authorization: `token ${useQueryCopilot.getState().notebookServerInfo.authToken}`, }, body: JSON.stringify(payload), }); @@ -215,13 +215,30 @@ export const QueryCopilotPromptbar: React.FC = ({ setGeneratedQueryComments(generateSQLQueryResponse.explanation); setShowFeedbackBar(true); resetQueryResults(); + TelemetryProcessor.traceSuccess(Action.QueryGenerationFromCopilotPrompt, { + databaseName: databaseId, + collectionId: containerId, + copilotLatency: + Date.parse(generateSQLQueryResponse?.generateEnd) - Date.parse(generateSQLQueryResponse?.generateStart), + responseCode: response.status, + }); } else { setShowInvalidQueryMessageBar(true); + TelemetryProcessor.traceFailure(Action.QueryGenerationFromCopilotPrompt, { + databaseName: databaseId, + collectionId: containerId, + responseCode: response.status, + }); } } else { handleError(JSON.stringify(generateSQLQueryResponse), "copilotInternalServerError"); useTabs.getState().setIsQueryErrorThrown(true); setShowErrorMessageBar(true); + TelemetryProcessor.traceFailure(Action.QueryGenerationFromCopilotPrompt, { + databaseName: databaseId, + collectionId: containerId, + responseCode: response.status, + }); } } catch (error) { handleError(error, "executeNaturalLanguageQuery"); @@ -310,6 +327,7 @@ export const QueryCopilotPromptbar: React.FC = ({ styles={{ root: { width: "95%" }, fieldGroup: { borderRadius: 6 } }} disabled={isGeneratingQuery} autoComplete="off" + placeholder="Ask a question in natural language and we’ll generate the query for you." /> {copilotTeachingBubbleVisible && ( = ({ description: "", userPrompt: userPrompt, }, - explorer: explorer, + explorer, + databaseId, + containerId, + mode: isSampleCopilotActive ? "Sample" : "User", }); }} directionalHint={DirectionalHint.topCenter} @@ -499,7 +520,7 @@ export const QueryCopilotPromptbar: React.FC = ({ { setShowCallout(false); - useQueryCopilot.getState().openFeedbackModal(generatedQuery, true, userPrompt); + openFeedbackModal(generatedQuery, true, userPrompt); }} > more feedback? @@ -524,7 +545,7 @@ export const QueryCopilotPromptbar: React.FC = ({ iconProps={{ iconName: dislikeQuery === true ? "DislikeSolid" : "Dislike" }} onClick={() => { if (!dislikeQuery) { - useQueryCopilot.getState().openFeedbackModal(generatedQuery, false, userPrompt); + openFeedbackModal(generatedQuery, false, userPrompt); setLikeQuery(false); } setDislikeQuery(!dislikeQuery); diff --git a/src/Explorer/QueryCopilot/QueryCopilotTab.tsx b/src/Explorer/QueryCopilot/QueryCopilotTab.tsx index 27d0559a2..c6e787ca0 100644 --- a/src/Explorer/QueryCopilot/QueryCopilotTab.tsx +++ b/src/Explorer/QueryCopilot/QueryCopilotTab.tsx @@ -1,5 +1,6 @@ /* eslint-disable no-console */ import { Stack } from "@fluentui/react"; +import { QueryCopilotSampleContainerId, QueryCopilotSampleDatabaseId } from "Common/Constants"; import { CommandButtonComponentProps } from "Explorer/Controls/CommandButton/CommandButtonComponent"; import { EditorReact } from "Explorer/Controls/Editor/EditorReact"; import { useCommandBar } from "Explorer/Menus/CommandBar/CommandBarComponentAdapter"; @@ -11,6 +12,7 @@ import { QueryCopilotResults } from "Explorer/QueryCopilot/Shared/QueryCopilotRe import { userContext } from "UserContext"; import { useQueryCopilot } from "hooks/useQueryCopilot"; import { useSidePanel } from "hooks/useSidePanel"; +import { ReactTabKind, TabsState, useTabs } from "hooks/useTabs"; import React, { useState } from "react"; import SplitterLayout from "react-splitter-layout"; import QueryCommandIcon from "../../../images/CopilotCommand.svg"; @@ -28,13 +30,14 @@ export const QueryCopilotTab: React.FC = ({ explorer }: Query ? StringUtility.toBoolean(cachedCopilotToggleStatus) : true; const [copilotActive, setCopilotActive] = useState(copilotInitialActive); + const [tabActive, setTabActive] = useState(true); const getCommandbarButtons = (): CommandButtonComponentProps[] => { const executeQueryBtnLabel = selectedQuery ? "Execute Selection" : "Execute Query"; const executeQueryBtn = { iconSrc: ExecuteQueryIcon, iconAlt: executeQueryBtnLabel, - onCommandClick: () => OnExecuteQueryClick(), + onCommandClick: () => OnExecuteQueryClick(useQueryCopilot), commandButtonLabel: executeQueryBtnLabel, ariaLabel: executeQueryBtnLabel, hasPopup: false, @@ -73,10 +76,13 @@ export const QueryCopilotTab: React.FC = ({ explorer }: Query React.useEffect(() => { return () => { - const commandbarButtons = getCommandbarButtons(); - commandbarButtons.pop(); - commandbarButtons.map((props: CommandButtonComponentProps) => (props.disabled = true)); - useCommandBar.getState().setContextButtons(commandbarButtons); + useTabs.subscribe((state: TabsState) => { + if (state.activeReactTab === ReactTabKind.QueryCopilot) { + setTabActive(true); + } else { + setTabActive(false); + } + }); }; }, []); @@ -88,8 +94,13 @@ export const QueryCopilotTab: React.FC = ({ explorer }: Query return (
- {copilotActive && ( - + {tabActive && copilotActive && ( + )} diff --git a/src/Explorer/QueryCopilot/QueryCopilotUtilities.ts b/src/Explorer/QueryCopilot/QueryCopilotUtilities.ts index f6d4392d9..7b0290f24 100644 --- a/src/Explorer/QueryCopilot/QueryCopilotUtilities.ts +++ b/src/Explorer/QueryCopilot/QueryCopilotUtilities.ts @@ -7,6 +7,11 @@ import { getCommonQueryOptions } from "Common/dataAccess/queryDocuments"; import DocumentId from "Explorer/Tree/DocumentId"; import { logConsoleProgress } from "Utils/NotificationConsoleUtils"; +export interface SuggestedPrompt { + id: number; + text: string; +} + export const querySampleDocuments = (query: string, options: FeedOptions): QueryIterator => { options = getCommonQueryOptions(options); return sampleDataClient() @@ -33,3 +38,19 @@ export const readSampleDocument = async (documentId: DocumentId): Promise clearMessage(); } }; + +export const getSampleDatabaseSuggestedPrompts = (): SuggestedPrompt[] => { + return [ + { id: 1, text: 'Show all products that have the word "ultra" in the name or description' }, + { id: 2, text: "What are all of the possible categories for the products, and their counts?" }, + { id: 3, text: 'Show me all products that have been reviewed by someone with a username that contains "bob"' }, + ]; +}; + +export const getSuggestedPrompts = (): SuggestedPrompt[] => { + return [ + { id: 1, text: "Show the first 10 items" }, + { id: 2, text: 'Count all the items in my data as "numItems"' }, + { id: 3, text: "Find the oldest item added to my collection" }, + ]; +}; diff --git a/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.test.ts b/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.test.ts index 40fa30b3a..8317f962e 100644 --- a/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.test.ts +++ b/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.test.ts @@ -1,4 +1,3 @@ -import { QueryCopilotSampleContainerSchema, ShortenedQueryCopilotSampleContainerSchema } from "Common/Constants"; import { handleError } from "Common/ErrorHandlingUtils"; import { createUri } from "Common/UrlUtility"; import Explorer from "Explorer/Explorer"; @@ -37,9 +36,6 @@ describe("Query Copilot Client", () => { userPrompt: "UserPrompt", description: "Description", contact: "Contact", - containerSchema: userContext.features.enableCopilotFullSchema - ? QueryCopilotSampleContainerSchema - : ShortenedQueryCopilotSampleContainerSchema, }; const mockStore = useQueryCopilot.getState(); @@ -59,6 +55,9 @@ describe("Query Copilot Client", () => { globalThis.fetch = mockFetch; await SubmitFeedback({ + databaseId: "test", + containerId: "test", + mode: "User", params: { likeQuery: true, generatedQuery: "GeneratedQuery", @@ -91,6 +90,9 @@ describe("Query Copilot Client", () => { globalThis.fetch = mockFetch; await SubmitFeedback({ + databaseId: "test", + containerId: "test", + mode: "User", params: { likeQuery: false, generatedQuery: "GeneratedQuery", @@ -108,6 +110,7 @@ describe("Query Copilot Client", () => { headers: { "content-type": "application/json", "x-ms-correlationid": "mocked-correlation-id", + Authorization: "token mocked-token", }, }), ); @@ -120,6 +123,9 @@ describe("Query Copilot Client", () => { globalThis.fetch = jest.fn().mockRejectedValueOnce(new Error("Mock error")); await SubmitFeedback({ + databaseId: "test", + containerId: "test", + mode: "User", params: { likeQuery: true, generatedQuery: "GeneratedQuery", diff --git a/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts b/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts index 9c9f1de61..b46cd9afc 100644 --- a/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts +++ b/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts @@ -1,28 +1,174 @@ import { FeedOptions } from "@azure/cosmos"; import { + Areas, + ConnectionStatusType, ContainerStatusType, + HttpStatusCodes, PoolIdType, QueryCopilotSampleContainerId, QueryCopilotSampleContainerSchema, ShortenedQueryCopilotSampleContainerSchema, } from "Common/Constants"; -import { getErrorMessage, handleError } from "Common/ErrorHandlingUtils"; +import { getErrorMessage, getErrorStack, handleError } from "Common/ErrorHandlingUtils"; import { shouldEnableCrossPartitionKey } from "Common/HeadersUtility"; import { MinimalQueryIterator } from "Common/IteratorUtilities"; import { createUri } from "Common/UrlUtility"; import { queryDocumentsPage } from "Common/dataAccess/queryDocumentsPage"; -import { QueryResults } from "Contracts/ViewModels"; +import { configContext } from "ConfigContext"; +import { + ContainerConnectionInfo, + CopilotEnabledConfiguration, + FeatureRegistration, + IProvisionData, +} from "Contracts/DataModels"; +import { AuthorizationTokenHeaderMetadata, QueryResults } from "Contracts/ViewModels"; +import { useDialog } from "Explorer/Controls/Dialog"; import Explorer from "Explorer/Explorer"; import { querySampleDocuments } from "Explorer/QueryCopilot/QueryCopilotUtilities"; import { FeedbackParams, GenerateSQLQueryResponse } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces"; import { Action } from "Shared/Telemetry/TelemetryConstants"; import { traceFailure, traceStart, traceSuccess } from "Shared/Telemetry/TelemetryProcessor"; import { userContext } from "UserContext"; +import { getAuthorizationHeader } from "Utils/AuthorizationUtils"; import { queryPagesUntilContentPresent } from "Utils/QueryUtils"; -import { useQueryCopilot } from "hooks/useQueryCopilot"; +import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot"; import { useTabs } from "hooks/useTabs"; import * as StringUtility from "../../../Shared/StringUtility"; +export const isCopilotFeatureRegistered = async (subscriptionId: string): Promise => { + const api_version = "2021-07-01"; + const url = `${configContext.ARM_ENDPOINT}/subscriptions/${subscriptionId}/providers/Microsoft.Features/featureProviders/Microsoft.DocumentDB/subscriptionFeatureRegistrations/CopilotInAzureCDB?api-version=${api_version}`; + const authorizationHeader: AuthorizationTokenHeaderMetadata = getAuthorizationHeader(); + const headers = { [authorizationHeader.header]: authorizationHeader.token }; + + let response; + + try { + response = await window.fetch(url, { + headers, + }); + } catch (error) { + return false; + } + + if (!response?.ok) { + return false; + } + + const featureRegistration = (await response?.json()) as FeatureRegistration; + return featureRegistration?.properties?.state === "Registered"; +}; + +export const getCopilotEnabled = async (): Promise => { + const url = `${configContext.BACKEND_ENDPOINT}/api/portalsettings/querycopilot`; + const authorizationHeader: AuthorizationTokenHeaderMetadata = getAuthorizationHeader(); + const headers = { [authorizationHeader.header]: authorizationHeader.token }; + + let response; + + try { + response = await window.fetch(url, { + headers, + }); + } catch (error) { + return false; + } + + if (!response?.ok) { + throw new Error(await response?.text()); + } + + const copilotPortalConfiguration = (await response?.json()) as CopilotEnabledConfiguration; + return copilotPortalConfiguration?.isEnabled; +}; + +export const allocatePhoenixContainer = async ({ + explorer, + databaseId, + containerId, + mode, +}: { + explorer: Explorer; + databaseId: string; + containerId: string; + mode: string; +}): Promise => { + try { + if ( + useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active && + !userContext.features.disableCopilotPhoenixGateaway + ) { + await explorer.allocateContainer(PoolIdType.QueryCopilot, mode); + } else { + const currentAllocatedSchemaInfo = useQueryCopilot.getState().schemaAllocationInfo; + if ( + currentAllocatedSchemaInfo.databaseId !== databaseId || + currentAllocatedSchemaInfo.containerId !== containerId + ) { + await resetPhoenixContainerSchema({ explorer, databaseId, containerId, mode }); + } + } + useQueryCopilot.getState().setSchemaAllocationInfo({ + databaseId, + containerId, + }); + } catch (error) { + traceFailure(Action.PhoenixConnection, { + dataExplorerArea: Areas.Copilot, + status: error.status, + error: getErrorMessage(error), + errorStack: getErrorStack(error), + }); + useQueryCopilot.getState().resetContainerConnection(); + if (error?.status === HttpStatusCodes.Forbidden && error.message) { + useDialog.getState().showOkModalDialog("Connection Failed", `${error.message}`); + } else { + useDialog + .getState() + .showOkModalDialog( + "Connection Failed", + "We are unable to connect to the temporary workspace. Please try again in a few minutes. If the error persists, file a support ticket.", + ); + } + } finally { + useTabs.getState().setIsTabExecuting(false); + } +}; + +export const resetPhoenixContainerSchema = async ({ + explorer, + databaseId, + containerId, + mode, +}: { + explorer: Explorer; + databaseId: string; + containerId: string; + mode: string; +}): Promise => { + try { + const provisionData: IProvisionData = { + poolId: PoolIdType.QueryCopilot, + databaseId: databaseId, + containerId: containerId, + mode: mode, + }; + const connectionInfo = await explorer.phoenixClient.allocateContainer(provisionData); + const connectionStatus: ContainerConnectionInfo = { + status: ConnectionStatusType.Connecting, + }; + await explorer.setNotebookInfo(false, connectionInfo, connectionStatus); + } catch (error) { + traceFailure(Action.PhoenixConnection, { + dataExplorerArea: Areas.Copilot, + status: error.status, + error: getErrorMessage(error), + errorStack: getErrorStack(error), + }); + throw error; + } +}; + export const SendQueryRequest = async ({ userPrompt, explorer, @@ -106,16 +252,19 @@ export const SendQueryRequest = async ({ export const SubmitFeedback = async ({ params, explorer, + databaseId, + containerId, + mode, }: { params: FeedbackParams; explorer: Explorer; + databaseId: string; + containerId: string; + mode: string; }): Promise => { try { const { likeQuery, generatedQuery, userPrompt, description, contact } = params; const payload = { - containerSchema: userContext.features.enableCopilotFullSchema - ? QueryCopilotSampleContainerSchema - : ShortenedQueryCopilotSampleContainerSchema, like: likeQuery ? "like" : "dislike", generatedSql: generatedQuery, userPrompt, @@ -126,7 +275,7 @@ export const SubmitFeedback = async ({ useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active && !userContext.features.disableCopilotPhoenixGateaway ) { - await explorer.allocateContainer(PoolIdType.QueryCopilot); + await allocatePhoenixContainer({ explorer, databaseId, containerId, mode }); } const serverInfo = useQueryCopilot.getState().notebookServerInfo; const feedbackUri = userContext.features.disableCopilotPhoenixGateaway @@ -137,6 +286,7 @@ export const SubmitFeedback = async ({ headers: { "content-type": "application/json", "x-ms-correlationid": useQueryCopilot.getState().correlationId, + Authorization: `token ${useQueryCopilot.getState().notebookServerInfo.authToken}`, }, body: JSON.stringify(payload), }); @@ -145,7 +295,7 @@ export const SubmitFeedback = async ({ } }; -export const OnExecuteQueryClick = async (): Promise => { +export const OnExecuteQueryClick = async (useQueryCopilot: Partial): Promise => { traceStart(Action.ExecuteQueryGeneratedFromQueryCopilot, { correlationId: useQueryCopilot.getState().correlationId, userPrompt: useQueryCopilot.getState().userPrompt, @@ -160,13 +310,14 @@ export const OnExecuteQueryClick = async (): Promise => { useQueryCopilot.getState().setQueryIterator(queryIterator); setTimeout(async () => { - await QueryDocumentsPerPage(0, queryIterator); + await QueryDocumentsPerPage(0, queryIterator, useQueryCopilot); }, 100); }; export const QueryDocumentsPerPage = async ( firstItemIndex: number, queryIterator: MinimalQueryIterator, + useQueryCopilot: Partial, ): Promise => { try { useQueryCopilot.getState().setIsExecuting(true); diff --git a/src/Explorer/QueryCopilot/Shared/QueryCopilotInterfaces.ts b/src/Explorer/QueryCopilot/Shared/QueryCopilotInterfaces.ts index fd22a7432..64d610a71 100644 --- a/src/Explorer/QueryCopilot/Shared/QueryCopilotInterfaces.ts +++ b/src/Explorer/QueryCopilot/Shared/QueryCopilotInterfaces.ts @@ -32,3 +32,8 @@ export interface FeedbackParams { export interface QueryCopilotProps { explorer: Explorer; } + +export interface CopilotSchemaAllocationInfo { + databaseId: string; + containerId: string; +} diff --git a/src/Explorer/QueryCopilot/Shared/QueryCopilotResults.tsx b/src/Explorer/QueryCopilot/Shared/QueryCopilotResults.tsx index 76bcc12e9..8d872e174 100644 --- a/src/Explorer/QueryCopilot/Shared/QueryCopilotResults.tsx +++ b/src/Explorer/QueryCopilot/Shared/QueryCopilotResults.tsx @@ -12,7 +12,7 @@ export const QueryCopilotResults: React.FC = (): JSX.Element => { queryResults={useQueryCopilot.getState().queryResults} isExecuting={useQueryCopilot.getState().isExecuting} executeQueryDocumentsPage={(firstItemIndex: number) => - QueryDocumentsPerPage(firstItemIndex, useQueryCopilot.getState().queryIterator) + QueryDocumentsPerPage(firstItemIndex, useQueryCopilot.getState().queryIterator, useQueryCopilot) } /> ); diff --git a/src/Explorer/QueryCopilot/__snapshots__/QueryCopilotTab.test.tsx.snap b/src/Explorer/QueryCopilot/__snapshots__/QueryCopilotTab.test.tsx.snap index 46713ab6c..26b52ff90 100644 --- a/src/Explorer/QueryCopilot/__snapshots__/QueryCopilotTab.test.tsx.snap +++ b/src/Explorer/QueryCopilot/__snapshots__/QueryCopilotTab.test.tsx.snap @@ -18,6 +18,8 @@ exports[`Query copilot tab snapshot test should render with initial input 1`] = } > { (state) => state.sampleDataResourceTokenCollection, ), }, + { + dispose: useQueryCopilot.subscribe( + () => this.setState({}), + (state) => state.copilotEnabled, + ), + }, ); } @@ -114,9 +121,9 @@ export class SplashScreen extends React.Component { private getSplashScreenButtons = (): JSX.Element => { if ( - useDatabases.getState().sampleDataResourceTokenCollection && - userContext.features.enableCopilot && - userContext.apiType === "SQL" + userContext.apiType === "SQL" && + useQueryCopilot.getState().copilotEnabled && + useDatabases.getState().sampleDataResourceTokenCollection ) { return ( diff --git a/src/Explorer/Tabs/DocumentsTab.ts b/src/Explorer/Tabs/DocumentsTab.ts index 8d95bbf18..d4e532ca4 100644 --- a/src/Explorer/Tabs/DocumentsTab.ts +++ b/src/Explorer/Tabs/DocumentsTab.ts @@ -1,10 +1,10 @@ import { ItemDefinition, PartitionKey, PartitionKeyDefinition, QueryIterator, Resource } from "@azure/cosmos"; import { querySampleDocuments, readSampleDocument } from "Explorer/QueryCopilot/QueryCopilotUtilities"; +import { QueryConstants } from "Shared/Constants"; +import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; import * as ko from "knockout"; import Q from "q"; import { format } from "react-string-format"; -import { QueryConstants } from "Shared/Constants"; -import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; import DeleteDocumentIcon from "../../../images/DeleteDocument.svg"; import NewDocumentIcon from "../../../images/NewDocument.svg"; import UploadIcon from "../../../images/Upload_16x16.svg"; @@ -331,6 +331,7 @@ export default class DocumentsTab extends TabsBase { this.showPartitionKey = this._shouldShowPartitionKey(); this._isQueryCopilotSampleContainer = + this.collection?.isSampleCollection && this.collection?.databaseId === QueryCopilotSampleDatabaseId && this.collection?.id() === QueryCopilotSampleContainerId; } diff --git a/src/Explorer/Tabs/QueryTab/QueryTab.tsx b/src/Explorer/Tabs/QueryTab/QueryTab.tsx index dba6f2f1e..47c2ce2ba 100644 --- a/src/Explorer/Tabs/QueryTab/QueryTab.tsx +++ b/src/Explorer/Tabs/QueryTab/QueryTab.tsx @@ -1,11 +1,16 @@ +import { CopilotProvider } from "Explorer/QueryCopilot/QueryCopilotContext"; +import { userContext } from "UserContext"; import React from "react"; import * as DataModels from "../../../Contracts/DataModels"; import type { QueryTabOptions } from "../../../Contracts/ViewModels"; import { useTabs } from "../../../hooks/useTabs"; import Explorer from "../../Explorer"; -import { IQueryTabComponentProps, ITabAccessor } from "../../Tabs/QueryTab/QueryTabComponent"; +import QueryTabComponent, { + IQueryTabComponentProps, + ITabAccessor, + QueryTabFunctionComponent, +} from "../../Tabs/QueryTab/QueryTabComponent"; import TabsBase from "../TabsBase"; -import QueryTabComponent from "./QueryTabComponent"; export interface IQueryTabProps { container: Explorer; @@ -40,7 +45,13 @@ export class NewQueryTab extends TabsBase { } public render(): JSX.Element { - return ; + return userContext.apiType === "SQL" ? ( + + + + ) : ( + + ); } public onTabClick(): void { diff --git a/src/Explorer/Tabs/QueryTab/QueryTabComponent.test.tsx b/src/Explorer/Tabs/QueryTab/QueryTabComponent.test.tsx index 36b9ac5f3..27774cfe0 100644 --- a/src/Explorer/Tabs/QueryTab/QueryTabComponent.test.tsx +++ b/src/Explorer/Tabs/QueryTab/QueryTabComponent.test.tsx @@ -1,6 +1,16 @@ import { fireEvent, render } from "@testing-library/react"; -import QueryTabComponent, { IQueryTabComponentProps } from "Explorer/Tabs/QueryTab/QueryTabComponent"; +import { CollectionTabKind } from "Contracts/ViewModels"; +import { CopilotProvider } from "Explorer/QueryCopilot/QueryCopilotContext"; +import { QueryCopilotPromptbar } from "Explorer/QueryCopilot/QueryCopilotPromptbar"; +import QueryTabComponent, { + IQueryTabComponentProps, + QueryTabFunctionComponent, +} from "Explorer/Tabs/QueryTab/QueryTabComponent"; +import TabsBase from "Explorer/Tabs/TabsBase"; +import { updateUserContext, userContext } from "UserContext"; +import { mount } from "enzyme"; import { useQueryCopilot } from "hooks/useQueryCopilot"; +import { useTabs } from "hooks/useTabs"; import React from "react"; jest.mock("Explorer/Controls/Editor/EditorReact"); @@ -11,9 +21,15 @@ describe("QueryTabComponent", () => { mockStore.showCopilotSidebar = false; mockStore.setShowCopilotSidebar = jest.fn(); }); - beforeEach(() => jest.clearAllMocks()); + afterEach(() => jest.clearAllMocks()); - it("should launch Copilot when ALT+C is pressed", () => { + it("should launch conversational Copilot when ALT+C is pressed and when copilot version is 3", () => { + updateUserContext({ + features: { + ...userContext.features, + copilotVersion: "v3.0", + }, + }); const propsMock: Readonly = { collection: { databaseId: "CopilotSampleDb" }, onTabAccessor: () => jest.fn(), @@ -31,4 +47,32 @@ describe("QueryTabComponent", () => { expect(mockStore.setShowCopilotSidebar).toHaveBeenCalledWith(true); }); + + it("copilot should be enabled by default when tab is active", () => { + useQueryCopilot.getState().setCopilotEnabled(true); + useQueryCopilot.getState().setCopilotUserDBEnabled(true); + const activeTab = new TabsBase({ + tabKind: CollectionTabKind.Query, + title: "Query", + tabPath: "", + }); + activeTab.tabId = "mockTabId"; + useTabs.getState().activeTab = activeTab; + const propsMock: Readonly = { + collection: { databaseId: "CopilotUserDb", id: () => "CopilotUserContainer" }, + onTabAccessor: () => jest.fn(), + isExecutionError: false, + tabId: "mockTabId", + tabsBaseInstance: { + updateNavbarWithTabsButtons: () => jest.fn(), + }, + } as unknown as IQueryTabComponentProps; + + const container = mount( + + + , + ); + expect(container.find(QueryCopilotPromptbar).exists()).toBe(true); + }); }); diff --git a/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx b/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx index d4deda184..4793f3791 100644 --- a/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx +++ b/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx @@ -1,21 +1,29 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +/* eslint-disable no-console */ import { FeedOptions } from "@azure/cosmos"; import { useDialog } from "Explorer/Controls/Dialog"; -import { OnExecuteQueryClick } from "Explorer/QueryCopilot/Shared/QueryCopilotClient"; -import { QueryCopilotResults } from "Explorer/QueryCopilot/Shared/QueryCopilotResults"; +import { QueryCopilotFeedbackModal } from "Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal"; +import { useCopilotStore } from "Explorer/QueryCopilot/QueryCopilotContext"; +import { QueryCopilotPromptbar } from "Explorer/QueryCopilot/QueryCopilotPromptbar"; +import { OnExecuteQueryClick, QueryDocumentsPerPage } from "Explorer/QueryCopilot/Shared/QueryCopilotClient"; import { QueryCopilotSidebar } from "Explorer/QueryCopilot/V2/Sidebar/QueryCopilotSidebar"; import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection"; +import { useSelectedNode } from "Explorer/useSelectedNode"; import { QueryConstants } from "Shared/Constants"; import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; +import { Action } from "Shared/Telemetry/TelemetryConstants"; import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot"; +import { TabsState, useTabs } from "hooks/useTabs"; import React, { Fragment } from "react"; import SplitterLayout from "react-splitter-layout"; import "react-splitter-layout/lib/index.css"; import { format } from "react-string-format"; +import QueryCommandIcon from "../../../../images/CopilotCommand.svg"; import LaunchCopilot from "../../../../images/CopilotTabIcon.svg"; import CancelQueryIcon from "../../../../images/Entity_cancel.svg"; import ExecuteQueryIcon from "../../../../images/ExecuteQuery.svg"; import SaveQueryIcon from "../../../../images/save-cosmos.svg"; -import { NormalizedEventKey, QueryCopilotSampleDatabaseId } from "../../../Common/Constants"; +import { NormalizedEventKey } from "../../../Common/Constants"; import { getErrorMessage } from "../../../Common/ErrorHandlingUtils"; import * as HeadersUtility from "../../../Common/HeadersUtility"; import { MinimalQueryIterator } from "../../../Common/IteratorUtilities"; @@ -24,6 +32,8 @@ import { queryDocuments } from "../../../Common/dataAccess/queryDocuments"; import { queryDocumentsPage } from "../../../Common/dataAccess/queryDocumentsPage"; import * as DataModels from "../../../Contracts/DataModels"; import * as ViewModels from "../../../Contracts/ViewModels"; +import * as StringUtility from "../../../Shared/StringUtility"; +import * as TelemetryProcessor from "../../../Shared/Telemetry/TelemetryProcessor"; import { userContext } from "../../../UserContext"; import * as QueryUtils from "../../../Utils/QueryUtils"; import { useSidePanel } from "../../../hooks/useSidePanel"; @@ -72,6 +82,9 @@ export interface IQueryTabComponentProps { isPreferredApiMongoDB?: boolean; monacoEditorSetting?: string; viewModelcollection?: ViewModels.Collection; + copilotEnabled?: boolean; + isSampleCopilotActive?: boolean; + copilotStore?: Partial; } interface IQueryTabStates { @@ -85,8 +98,24 @@ interface IQueryTabStates { showCopilotSidebar: boolean; queryCopilotGeneratedQuery: string; cancelQueryTimeoutID: NodeJS.Timeout; + copilotActive: boolean; + currentTabActive: boolean; } +export const QueryTabFunctionComponent = (props: IQueryTabComponentProps): any => { + const copilotStore = useCopilotStore(); + const isSampleCopilotActive = useSelectedNode.getState().isQueryCopilotCollectionSelected(); + const queryTabProps = { + ...props, + copilotEnabled: + useQueryCopilot().copilotEnabled && + (useQueryCopilot().copilotUserDBEnabled || (isSampleCopilotActive && !!userContext.sampleDataConnectionInfo)), + isSampleCopilotActive: isSampleCopilotActive, + copilotStore: copilotStore, + }; + return ; +}; + export default class QueryTabComponent extends React.Component { public queryEditorId: string; public executeQueryButton: Button; @@ -113,12 +142,14 @@ export default class QueryTabComponent extends React.Component 0, visible: true, @@ -143,6 +174,19 @@ export default class QueryTabComponent extends React.Component { await this._executeQueryDocumentsPage(0); }, 100); + if (this.state.copilotActive) { + const query = this.state.sqlQueryEditorContent.split("\r\n")?.pop(); + const isqueryEdited = this.props.copilotStore.generatedQuery && this.props.copilotStore.generatedQuery !== query; + if (isqueryEdited) { + TelemetryProcessor.traceMark(Action.QueryEdited, { + databaseName: this.props.collection.databaseId, + collectionId: this.props.collection.id(), + }); + } + } }; public onSaveQueryClick = (): void => { @@ -326,7 +380,9 @@ export default class QueryTabComponent extends React.Component OnExecuteQueryClick() : this.onExecuteQueryClick, + onCommandClick: this.props.isSampleCopilotActive + ? () => OnExecuteQueryClick(this.props.copilotStore) + : this.onExecuteQueryClick, commandButtonLabel: label, ariaLabel: label, hasPopup: false, @@ -380,6 +436,20 @@ export default class QueryTabComponent extends React.Component { + this._toggleCopilot(!this.state.copilotActive); + }, + commandButtonLabel: "Copilot", + ariaLabel: "Copilot", + hasPopup: false, + }; + buttons.push(toggleCopilotButton); + } + if (!this.props.isPreferredApiMongoDB && this.state.isExecuting) { const label = "Cancel query"; buttons.push({ @@ -395,11 +465,31 @@ export default class QueryTabComponent extends React.Component { + this.setState({ copilotActive: active }); + useQueryCopilot.getState().setCopilotEnabledforExecution(active); + localStorage.setItem(`${userContext.databaseAccount?.id}-queryCopilotToggleStatus`, active.toString()); + + TelemetryProcessor.traceSuccess(active ? Action.ActivateQueryCopilot : Action.DeactivateQueryCopilot, { + databaseName: this.props.collection.databaseId, + collectionId: this.props.collection.id(), + }); + }; + + componentDidUpdate = (_prevProps: IQueryTabComponentProps, prevState: IQueryTabStates): void => { + if (prevState.copilotActive !== this.state.copilotActive) { + useCommandBar.getState().setContextButtons(this.getTabsButtons()); + } + }; + public onChangeContent(newContent: string): void { this.setState({ sqlQueryEditorContent: newContent, queryCopilotGeneratedQuery: "", }); + if (this.state.copilotActive) { + this.props.copilotStore?.setQuery(newContent); + } if (this.isPreferredApiMongoDB) { if (newContent.length > 0) { this.executeQueryButton = { @@ -434,6 +524,10 @@ export default class QueryTabComponent extends React.Component void; componentDidMount(): void { - this.unsubscribeCopilotSidebar = useQueryCopilot.subscribe((state: QueryCopilotState) => { - if (this.state.showCopilotSidebar !== state.showCopilotSidebar) { - this.setState({ showCopilotSidebar: state.showCopilotSidebar }); - } - if (this.state.queryCopilotGeneratedQuery !== state.query) { - this.setState({ queryCopilotGeneratedQuery: state.query }); + useTabs.subscribe((state: TabsState) => { + if (this.state.currentTabActive && state.activeTab?.tabId !== this.props.tabId) { + this.setState({ + currentTabActive: false, + }); + } else if (!this.state.currentTabActive && state.activeTab?.tabId === this.props.tabId) { + this.setState({ + currentTabActive: true, + }); } }); @@ -466,7 +567,6 @@ export default class QueryTabComponent extends React.Component
+ {this.props.copilotEnabled && this.state.currentTabActive && this.state.copilotActive && ( + + )}
@@ -482,6 +590,7 @@ export default class QueryTabComponent extends React.Component this.onChangeContent(newContent)} @@ -489,8 +598,21 @@ export default class QueryTabComponent extends React.Component
- {this.isCopilotTabActive ? ( - + {this.props.isSampleCopilotActive ? ( + + QueryDocumentsPerPage( + firstItemIndex, + this.props.copilotStore.queryIterator, + this.props.copilotStore, + ) + } + /> ) : (
+ {this.props.copilotEnabled && this.props.copilotStore?.showFeedbackModal && ( + + )} ); } diff --git a/src/Explorer/Tabs/Tabs.tsx b/src/Explorer/Tabs/Tabs.tsx index d731bdf58..923b356a2 100644 --- a/src/Explorer/Tabs/Tabs.tsx +++ b/src/Explorer/Tabs/Tabs.tsx @@ -11,7 +11,6 @@ import { QuickstartTab } from "Explorer/Tabs/QuickstartTab"; import { VcoreMongoConnectTab } from "Explorer/Tabs/VCoreMongoConnectTab"; import { VcoreMongoQuickstartTab } from "Explorer/Tabs/VCoreMongoQuickstartTab"; import { userContext } from "UserContext"; -import { useQueryCopilot } from "hooks/useQueryCopilot"; import { useTeachingBubble } from "hooks/useTeachingBubble"; import ko from "knockout"; import React, { MutableRefObject, useEffect, useRef, useState } from "react"; @@ -170,7 +169,7 @@ const CloseButton = ({ onClick={(event: React.MouseEvent) => { event.stopPropagation(); tab ? tab.onCloseTabButtonClick() : useTabs.getState().closeReactTab(tabKind); - tabKind === ReactTabKind.QueryCopilot && useQueryCopilot.getState().resetQueryCopilotStates(); + // tabKind === ReactTabKind.QueryCopilot && useQueryCopilot.getState().resetQueryCopilotStates(); }} tabIndex={active ? 0 : undefined} onKeyPress={({ nativeEvent: e }) => tab.onKeyPressClose(undefined, e)} @@ -256,6 +255,7 @@ const isQueryErrorThrown = (tab?: Tab, tabKind?: ReactTabKind): boolean => { }; const getReactTabContent = (activeReactTab: ReactTabKind, explorer: Explorer): JSX.Element => { + // eslint-disable-next-line no-console switch (activeReactTab) { case ReactTabKind.Connect: return userContext.apiType === "VCoreMongo" ? ( diff --git a/src/Explorer/Tabs/TabsBase.ts b/src/Explorer/Tabs/TabsBase.ts index 308bb451f..a6f3a45dd 100644 --- a/src/Explorer/Tabs/TabsBase.ts +++ b/src/Explorer/Tabs/TabsBase.ts @@ -3,15 +3,15 @@ import * as Constants from "../../Common/Constants"; import * as ThemeUtility from "../../Common/ThemeUtility"; import * as DataModels from "../../Contracts/DataModels"; import * as ViewModels from "../../Contracts/ViewModels"; -import { useNotificationConsole } from "../../hooks/useNotificationConsole"; -import { useTabs } from "../../hooks/useTabs"; import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstants"; import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor"; +import { useNotificationConsole } from "../../hooks/useNotificationConsole"; +import { useTabs } from "../../hooks/useTabs"; import { CommandButtonComponentProps } from "../Controls/CommandButton/CommandButtonComponent"; import Explorer from "../Explorer"; import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter"; -import { useSelectedNode } from "../useSelectedNode"; import { WaitsForTemplateViewModel } from "../WaitsForTemplateViewModel"; +import { useSelectedNode } from "../useSelectedNode"; // TODO: Use specific actions for logging telemetry data export default class TabsBase extends WaitsForTemplateViewModel { private static id = 0; diff --git a/src/Explorer/Tree/Collection.ts b/src/Explorer/Tree/Collection.ts index 214456a63..4b5411d09 100644 --- a/src/Explorer/Tree/Collection.ts +++ b/src/Explorer/Tree/Collection.ts @@ -177,8 +177,8 @@ export default class Collection implements ViewModels.Collection { this.children.subscribe(() => { // update the database in zustand store const database = this.getDatabase(); - database.collections( - database.collections()?.map((collection) => { + database?.collections( + database?.collections()?.map((collection) => { if (collection.id() === this.id()) { return this; } diff --git a/src/Explorer/Tree/ResourceTokenCollection.ts b/src/Explorer/Tree/ResourceTokenCollection.ts index 4d2939824..a5e6ab07f 100644 --- a/src/Explorer/Tree/ResourceTokenCollection.ts +++ b/src/Explorer/Tree/ResourceTokenCollection.ts @@ -156,6 +156,6 @@ export default class ResourceTokenCollection implements ViewModels.CollectionBas } public getDatabase(): ViewModels.Database { - return useDatabases.getState().findDatabaseWithId(this.databaseId); + return useDatabases.getState().findDatabaseWithId(this.databaseId, this.isSampleCollection); } } diff --git a/src/Explorer/Tree/ResourceTree.tsx b/src/Explorer/Tree/ResourceTree.tsx index 27d815fdd..8eaaa9e66 100644 --- a/src/Explorer/Tree/ResourceTree.tsx +++ b/src/Explorer/Tree/ResourceTree.tsx @@ -1,6 +1,7 @@ import { Callout, DirectionalHint, ICalloutProps, ILinkProps, Link, Stack, Text } from "@fluentui/react"; import { SampleDataTree } from "Explorer/Tree/SampleDataTree"; import { getItemName } from "Utils/APITypeUtils"; +import { useQueryCopilot } from "hooks/useQueryCopilot"; import * as React from "react"; import shallow from "zustand/shallow"; import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg"; @@ -767,7 +768,8 @@ export const ResourceTree: React.FC = ({ container }: Resourc }; const dataRootNode = buildDataTree(); - const isSampleDataEnabled = userContext.sampleDataConnectionInfo && userContext.apiType === "SQL"; + const isSampleDataEnabled = + useQueryCopilot().copilotEnabled && userContext.sampleDataConnectionInfo && userContext.apiType === "SQL"; const sampleDataResourceTokenCollection = useDatabases((state) => state.sampleDataResourceTokenCollection); return ( diff --git a/src/Explorer/useDatabases.ts b/src/Explorer/useDatabases.ts index f47a5fb85..7701f1360 100644 --- a/src/Explorer/useDatabases.ts +++ b/src/Explorer/useDatabases.ts @@ -14,7 +14,7 @@ interface DatabasesState { deleteDatabase: (database: ViewModels.Database) => void; clearDatabases: () => void; isSaveQueryEnabled: () => boolean; - findDatabaseWithId: (databaseId: string) => ViewModels.Database; + findDatabaseWithId: (databaseId: string, isSampleDatabase?: boolean) => ViewModels.Database; isLastNonEmptyDatabase: () => boolean; findCollection: (databaseId: string, collectionId: string) => ViewModels.Collection; isLastCollection: () => boolean; @@ -33,7 +33,7 @@ export const useDatabases: UseStore = create((set, get) => ({ updateDatabase: (updatedDatabase: ViewModels.Database) => set((state) => { const updatedDatabases = state.databases.map((database: ViewModels.Database) => { - if (database.id() === updatedDatabase.id()) { + if (database?.id() === updatedDatabase?.id()) { return updatedDatabase; } @@ -67,7 +67,9 @@ export const useDatabases: UseStore = create((set, get) => ({ } return true; }, - findDatabaseWithId: (databaseId: string) => get().databases.find((db) => databaseId === db.id()), + findDatabaseWithId: (databaseId: string, isSampleDatabase?: boolean) => { + return get().databases.find((db) => databaseId === db.id() && db.isSampleDB === isSampleDatabase); + }, isLastNonEmptyDatabase: () => { const databases = get().databases; return databases.length === 1 && (databases[0].collections()?.length > 0 || !!databases[0].offer()); diff --git a/src/Main.tsx b/src/Main.tsx index a76f314bc..d62f9b45e 100644 --- a/src/Main.tsx +++ b/src/Main.tsx @@ -1,13 +1,13 @@ // CSS Dependencies import { initializeIcons, loadTheme } from "@fluentui/react"; -import "bootstrap/dist/css/bootstrap.css"; import { QuickstartCarousel } from "Explorer/Quickstart/QuickstartCarousel"; import { MongoQuickstartTutorial } from "Explorer/Quickstart/Tutorials/MongoQuickstartTutorial"; import { SQLQuickstartTutorial } from "Explorer/Quickstart/Tutorials/SQLQuickstartTutorial"; +import { userContext } from "UserContext"; +import "bootstrap/dist/css/bootstrap.css"; import { useCarousel } from "hooks/useCarousel"; import React, { useState } from "react"; import ReactDOM from "react-dom"; -import { userContext } from "UserContext"; import "../externals/jquery-ui.min.css"; import "../externals/jquery-ui.min.js"; import "../externals/jquery-ui.structure.min.css"; @@ -18,21 +18,19 @@ import "../externals/jquery.typeahead.min.js"; // Image Dependencies import { Platform } from "ConfigContext"; import { QueryCopilotCarousel } from "Explorer/QueryCopilot/CopilotCarousel"; -import { QueryCopilotFeedbackModal } from "Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal"; -import { useQueryCopilot } from "hooks/useQueryCopilot"; import "../images/CosmosDB_rgb_ui_lighttheme.ico"; -import "../images/favicon.ico"; import hdeConnectImage from "../images/HdeConnectCosmosDB.svg"; +import "../images/favicon.ico"; +import "../less/TableStyles/CustomizeColumns.less"; +import "../less/TableStyles/EntityEditor.less"; +import "../less/TableStyles/fulldatatables.less"; +import "../less/TableStyles/queryBuilder.less"; import "../less/documentDB.less"; import "../less/forms.less"; import "../less/infobox.less"; import "../less/menus.less"; import "../less/messagebox.less"; import "../less/resourceTree.less"; -import "../less/TableStyles/CustomizeColumns.less"; -import "../less/TableStyles/EntityEditor.less"; -import "../less/TableStyles/fulldatatables.less"; -import "../less/TableStyles/queryBuilder.less"; import "../less/tree.less"; import { CollapsedResourceTree } from "./Common/CollapsedResourceTree"; import { ResourceTreeContainer } from "./Common/ResourceTreeContainer"; @@ -55,11 +53,11 @@ import "./Explorer/Panes/PanelComponent.less"; import { SidePanel } from "./Explorer/Panes/PanelContainerComponent"; import "./Explorer/SplashScreen/SplashScreen.less"; import { Tabs } from "./Explorer/Tabs/Tabs"; -import { useConfig } from "./hooks/useConfig"; -import { useKnockoutExplorer } from "./hooks/useKnockoutExplorer"; import "./Libs/jquery"; import { appThemeFabric } from "./Platform/Fabric/FabricTheme"; import "./Shared/appInsights"; +import { useConfig } from "./hooks/useConfig"; +import { useKnockoutExplorer } from "./hooks/useKnockoutExplorer"; initializeIcons(); @@ -67,7 +65,6 @@ const App: React.FunctionComponent = () => { const [isLeftPaneExpanded, setIsLeftPaneExpanded] = useState(true); const isCarouselOpen = useCarousel((state) => state.shouldOpen); const isCopilotCarouselOpen = useCarousel((state) => state.showCopilotCarousel); - const shouldShowModal = useQueryCopilot((state) => state.showFeedbackModal); const config = useConfig(); if (config?.platform === Platform.Fabric) { @@ -136,7 +133,6 @@ const App: React.FunctionComponent = () => { {} {} {} - {shouldShowModal && } ); }; diff --git a/src/Platform/Hosted/extractFeatures.ts b/src/Platform/Hosted/extractFeatures.ts index ef04eff14..7bf3c8a3f 100644 --- a/src/Platform/Hosted/extractFeatures.ts +++ b/src/Platform/Hosted/extractFeatures.ts @@ -112,7 +112,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear enableLegacyMongoShellV2Debug: "true" === get("enablelegacymongoshellv2debug"), loadLegacyMongoShellFromBE: "true" === get("loadlegacymongoshellfrombe"), enableCopilot: "true" === get("enablecopilot", "true"), - copilotVersion: get("copilotversion") ?? "v1.0", + copilotVersion: get("copilotversion") ?? "v2.0", disableCopilotPhoenixGateaway: "true" === get("disablecopilotphoenixgateaway"), enableCopilotFullSchema: "true" === get("enablecopilotfullschema", "true"), copilotChatFixedMonacoEditorHeight: "true" === get("copilotchatfixedmonacoeditorheight"), diff --git a/src/Shared/Telemetry/TelemetryConstants.ts b/src/Shared/Telemetry/TelemetryConstants.ts index d9d973dc4..c6211d3f9 100644 --- a/src/Shared/Telemetry/TelemetryConstants.ts +++ b/src/Shared/Telemetry/TelemetryConstants.ts @@ -133,6 +133,10 @@ export enum Action { CompleteUITour, OpenQueryCopilotFromSplashScreen, OpenQueryCopilotFromNewQuery, + ActivateQueryCopilot, + DeactivateQueryCopilot, + QueryGenerationFromCopilotPrompt, + QueryEdited, ExecuteQueryGeneratedFromQueryCopilot, } diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts index f515eb2d3..24aca7295 100644 --- a/src/hooks/useKnockoutExplorer.ts +++ b/src/hooks/useKnockoutExplorer.ts @@ -7,6 +7,7 @@ import { scheduleRefreshDatabaseResourceToken, } from "Platform/Fabric/FabricUtil"; import { getNetworkSettingsWarningMessage } from "Utils/NetworkUtility"; +import { useQueryCopilot } from "hooks/useQueryCopilot"; import { ReactTabKind, useTabs } from "hooks/useTabs"; import { useEffect, useState } from "react"; import { AuthType } from "../AuthType"; @@ -79,7 +80,7 @@ export function useKnockoutExplorer(platform: Platform): Explorer { if (explorer) { applyExplorerBindings(explorer); if (userContext.features.enableCopilot) { - updateContextForSampleData(explorer); + updateContextForCopilot(explorer).then(() => updateContextForSampleData(explorer)); } } }, [explorer]); @@ -554,12 +555,23 @@ interface PortalMessage { inputs?: DataExplorerInputsFrame; } +async function updateContextForCopilot(explorer: Explorer): Promise { + await explorer.configureCopilot(); +} + async function updateContextForSampleData(explorer: Explorer): Promise { - if (!userContext.features.enableCopilot) { + const copilotEnabled = + userContext.apiType === "SQL" && userContext.features.enableCopilot && useQueryCopilot.getState().copilotEnabled; + + if (!copilotEnabled) { return; } - const url = createUri(`${configContext.BACKEND_ENDPOINT}`, `/api/tokens/sampledataconnection`); + const sampleDatabaseEndpoint = useQueryCopilot.getState().copilotUserDBEnabled + ? `/api/tokens/sampledataconnection/v2` + : `/api/tokens/sampledataconnection`; + + const url = createUri(`${configContext.BACKEND_ENDPOINT}`, sampleDatabaseEndpoint); const authorizationHeader = getAuthorizationHeader(); const headers = { [authorizationHeader.header]: authorizationHeader.token }; diff --git a/src/hooks/useQueryCopilot.ts b/src/hooks/useQueryCopilot.ts index 9047c1853..5bcbfb861 100644 --- a/src/hooks/useQueryCopilot.ts +++ b/src/hooks/useQueryCopilot.ts @@ -1,6 +1,6 @@ import { MinimalQueryIterator } from "Common/IteratorUtilities"; import { QueryResults } from "Contracts/ViewModels"; -import { CopilotMessage } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces"; +import { CopilotMessage, CopilotSchemaAllocationInfo } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces"; import { guid } from "Explorer/Tables/Utilities"; import { useTabs } from "hooks/useTabs"; import create, { UseStore } from "zustand"; @@ -8,6 +8,8 @@ import * as DataModels from "../Contracts/DataModels"; import { ContainerInfo } from "../Contracts/DataModels"; export interface QueryCopilotState { + copilotEnabled: boolean; + copilotUserDBEnabled: boolean; generatedQuery: string; likeQuery: boolean; userPrompt: string; @@ -40,8 +42,14 @@ export interface QueryCopilotState { showExplanationBubble: boolean; notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo; containerStatus: ContainerInfo; + schemaAllocationInfo: CopilotSchemaAllocationInfo; isAllocatingContainer: boolean; + copilotEnabledforExecution: boolean; + getState?: () => QueryCopilotState; + + setCopilotEnabled: (copilotEnabled: boolean) => void; + setCopilotUserDBEnabled: (copilotUserDBEnabled: boolean) => void; openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void; closeFeedbackModal: () => void; setHideFeedbackModalForLikedQueries: (hideFeedbackModalForLikedQueries: boolean) => void; @@ -76,6 +84,8 @@ export interface QueryCopilotState { setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => void; setContainerStatus: (containerStatus: ContainerInfo) => void; setIsAllocatingContainer: (isAllocatingContainer: boolean) => void; + setSchemaAllocationInfo: (schemaAllocationInfo: CopilotSchemaAllocationInfo) => void; + setCopilotEnabledforExecution: (copilotEnabledforExecution: boolean) => void; resetContainerConnection: () => void; resetQueryCopilotStates: () => void; @@ -84,6 +94,8 @@ export interface QueryCopilotState { type QueryCopilotStore = UseStore; export const useQueryCopilot: QueryCopilotStore = create((set) => ({ + copilotEnabled: false, + copilotUserDBEnabled: false, generatedQuery: "", likeQuery: false, userPrompt: "", @@ -124,8 +136,15 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({ durationLeftInMinutes: undefined, phoenixServerInfo: undefined, }, + schemaAllocationInfo: { + databaseId: undefined, + containerId: undefined, + }, isAllocatingContainer: false, + copilotEnabledforExecution: false, + setCopilotEnabled: (copilotEnabled: boolean) => set({ copilotEnabled }), + setCopilotUserDBEnabled: (copilotUserDBEnabled: boolean) => set({ copilotUserDBEnabled }), openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }), closeFeedbackModal: () => set({ showFeedbackModal: false }), @@ -163,6 +182,8 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({ set({ notebookServerInfo }), setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }), setIsAllocatingContainer: (isAllocatingContainer: boolean) => set({ isAllocatingContainer }), + setSchemaAllocationInfo: (schemaAllocationInfo: CopilotSchemaAllocationInfo) => set({ schemaAllocationInfo }), + setCopilotEnabledforExecution: (copilotEnabledforExecution: boolean) => set({ copilotEnabledforExecution }), resetContainerConnection: (): void => { useTabs.getState().closeAllNotebookTabs(true); @@ -173,6 +194,10 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({ durationLeftInMinutes: undefined, phoenixServerInfo: undefined, }); + useQueryCopilot.getState().setSchemaAllocationInfo({ + databaseId: undefined, + containerId: undefined, + }); }, resetQueryCopilotStates: () => { @@ -217,6 +242,10 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({ durationLeftInMinutes: undefined, phoenixServerInfo: undefined, }, + schemaAllocationInfo: { + databaseId: undefined, + containerId: undefined, + }, isAllocatingContainer: false, })); }, diff --git a/src/hooks/useTabs.ts b/src/hooks/useTabs.ts index a1f3f4f2b..136e0ce2c 100644 --- a/src/hooks/useTabs.ts +++ b/src/hooks/useTabs.ts @@ -5,7 +5,7 @@ import NotebookTabV2 from "../Explorer/Tabs/NotebookV2Tab"; import TabsBase from "../Explorer/Tabs/TabsBase"; import { Platform, configContext } from "./../ConfigContext"; -interface TabsState { +export interface TabsState { openedTabs: TabsBase[]; openedReactTabs: ReactTabKind[]; activeTab: TabsBase | undefined;