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
2025-12-21 09:51:11 +00:00 · 2023-11-09 11:55:25 -06:00
parent a5e1b37ba6
commit 0e124f4881
38 changed files with 911 additions and 252 deletions
--- 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<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>
    );
  }