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
This commit is contained in:
sunghyunkang1111
2023-11-09 11:55:25 -06:00
committed by GitHub
parent a5e1b37ba6
commit 0e124f4881
38 changed files with 911 additions and 252 deletions

View File

@@ -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 <QueryTabComponent {...this.iQueryTabComponentProps} />;
return userContext.apiType === "SQL" ? (
<CopilotProvider>
<QueryTabFunctionComponent {...this.iQueryTabComponentProps} />
</CopilotProvider>
) : (
<QueryTabComponent {...this.iQueryTabComponentProps} />
);
}
public onTabClick(): void {

View File

@@ -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<IQueryTabComponentProps> = {
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<IQueryTabComponentProps> = {
collection: { databaseId: "CopilotUserDb", id: () => "CopilotUserContainer" },
onTabAccessor: () => jest.fn(),
isExecutionError: false,
tabId: "mockTabId",
tabsBaseInstance: {
updateNavbarWithTabsButtons: () => jest.fn(),
},
} as unknown as IQueryTabComponentProps;
const container = mount(
<CopilotProvider>
<QueryTabFunctionComponent {...propsMock} />
</CopilotProvider>,
);
expect(container.find(QueryCopilotPromptbar).exists()).toBe(true);
});
});

View File

@@ -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<QueryCopilotState>;
}
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 <QueryTabComponent {...queryTabProps}></QueryTabComponent>;
};
export default class QueryTabComponent extends React.Component<IQueryTabComponentProps, IQueryTabStates> {
public queryEditorId: string;
public executeQueryButton: Button;
@@ -113,12 +142,14 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
showCopilotSidebar: useQueryCopilot.getState().showCopilotSidebar,
queryCopilotGeneratedQuery: useQueryCopilot.getState().query,
cancelQueryTimeoutID: undefined,
copilotActive: this._queryCopilotActive(),
currentTabActive: true,
};
this.isCloseClicked = false;
this.splitterId = this.props.tabId + "_splitter";
this.queryEditorId = `queryeditor${this.props.tabId}`;
this.isPreferredApiMongoDB = this.props.isPreferredApiMongoDB;
this.isCopilotTabActive = QueryCopilotSampleDatabaseId === this.props.collection.databaseId;
this.isCopilotTabActive = userContext.features.copilotVersion === "v3.0";
this.executeQueryButton = {
enabled: !!this.state.sqlQueryEditorContent && this.state.sqlQueryEditorContent.length > 0,
visible: true,
@@ -143,6 +174,19 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
});
}
private _queryCopilotActive(): boolean {
if (this.props.copilotEnabled) {
const cachedCopilotToggleStatus: string = localStorage.getItem(
`${userContext.databaseAccount?.id}-queryCopilotToggleStatus`,
);
const copilotInitialActive: boolean = cachedCopilotToggleStatus
? StringUtility.toBoolean(cachedCopilotToggleStatus)
: true;
return copilotInitialActive;
}
return false;
}
public onCloseClick(isClicked: boolean): void {
this.isCloseClicked = isClicked;
if (useQueryCopilot.getState().wasCopilotUsed && this.isCopilotTabActive) {
@@ -167,6 +211,16 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
setTimeout(async () => {
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<IQueryTabComponen
buttons.push({
iconSrc: ExecuteQueryIcon,
iconAlt: label,
onCommandClick: this.isCopilotTabActive ? () => 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<IQueryTabComponen
buttons.push(launchCopilotButton);
}
if (this.props.copilotEnabled) {
const toggleCopilotButton = {
iconSrc: QueryCommandIcon,
iconAlt: "Copilot",
onCommandClick: () => {
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<IQueryTabComponen
return buttons;
}
private _toggleCopilot = (active: boolean) => {
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<IQueryTabComponen
: useQueryCopilot.getState().setSelectedQuery("");
}
if (this.state.copilotActive) {
this.props.copilotStore?.setSelectedQuery(selectedContent);
}
useCommandBar.getState().setContextButtons(this.getTabsButtons());
}
@@ -442,6 +536,10 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
return this.state.queryCopilotGeneratedQuery;
}
if (this.state.copilotActive) {
return this.props.copilotStore?.query;
}
return this.state.sqlQueryEditorContent;
}
@@ -452,12 +550,15 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
private unsubscribeCopilotSidebar: () => 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<IQueryTabComponen
}
componentWillUnmount(): void {
this.unsubscribeCopilotSidebar();
document.removeEventListener("keydown", this.handleCopilotKeyDown);
}
@@ -474,6 +574,14 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
return (
<Fragment>
<div className="tab-pane" id={this.props.tabId} role="tabpanel">
{this.props.copilotEnabled && this.state.currentTabActive && this.state.copilotActive && (
<QueryCopilotPromptbar
explorer={this.props.collection.container}
toggleCopilot={this._toggleCopilot}
databaseId={this.props.collection.databaseId}
containerId={this.props.collection.id()}
></QueryCopilotPromptbar>
)}
<div className="tabPaneContentContainer">
<SplitterLayout vertical={true} primaryIndex={0} primaryMinSize={100} secondaryMinSize={200}>
<Fragment>
@@ -482,6 +590,7 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
language={"sql"}
content={this.setEditorContent()}
isReadOnly={false}
wordWrap={"on"}
ariaLabel={"Editing Query"}
lineNumbers={"on"}
onContentChanged={(newContent: string) => this.onChangeContent(newContent)}
@@ -489,8 +598,21 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
/>
</div>
</Fragment>
{this.isCopilotTabActive ? (
<QueryCopilotResults />
{this.props.isSampleCopilotActive ? (
<QueryResultSection
isMongoDB={this.props.isPreferredApiMongoDB}
queryEditorContent={this.state.sqlQueryEditorContent}
error={this.props.copilotStore?.errorMessage}
queryResults={this.props.copilotStore?.queryResults}
isExecuting={this.props.copilotStore?.isExecuting}
executeQueryDocumentsPage={(firstItemIndex: number) =>
QueryDocumentsPerPage(
firstItemIndex,
this.props.copilotStore.queryIterator,
this.props.copilotStore,
)
}
/>
) : (
<QueryResultSection
isMongoDB={this.props.isPreferredApiMongoDB}
@@ -506,6 +628,14 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
</SplitterLayout>
</div>
</div>
{this.props.copilotEnabled && this.props.copilotStore?.showFeedbackModal && (
<QueryCopilotFeedbackModal
explorer={this.props.collection.container}
databaseId={this.props.collection.databaseId}
containerId={this.props.collection.id()}
mode={this.props.isSampleCopilotActive ? "Sample" : "User"}
/>
)}
</Fragment>
);
}