[Query Copilot] Add telemetry for query copilot (#1555)

src/Explorer/ContextMenuButtonFactory.tsx

@@ -1,3 +1,5 @@
+import { Action } from "Shared/Telemetry/TelemetryConstants";
+import { traceOpen } from "Shared/Telemetry/TelemetryProcessor";
 import { ReactTabKind, useTabs } from "hooks/useTabs";
 import React from "react";
 import AddCollectionIcon from "../../images/AddCollection.svg";
@@ -146,7 +148,10 @@ export const createSampleCollectionContextMenuButton = (): TreeNodeMenuItem[] =>
   if (userContext.apiType === "SQL") {
     items.push({
       iconSrc: AddSqlQueryIcon,
-      onClick: () => useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot),
+      onClick: () => {
+        useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot);
+        traceOpen(Action.OpenQueryCopilotFromNewQuery, { apiType: userContext.apiType });
+      },
       label: "New SQL Query",
     });
   }

src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx

@@ -1,3 +1,5 @@
+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";
@@ -326,6 +328,7 @@ function createNewSQLQueryButton(selectedNodeState: SelectedNodeState): CommandB
       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);

src/Explorer/QueryCopilot/QueryCopilotTab.tsx

@@ -34,6 +34,8 @@ import { DeletePopup } from "Explorer/QueryCopilot/Popup/DeletePopup";
 import { querySampleDocuments, submitFeedback } from "Explorer/QueryCopilot/QueryCopilotUtilities";
 import { SamplePrompts, SamplePromptsProps } from "Explorer/QueryCopilot/SamplePrompts/SamplePrompts";
 import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection";
+import { Action } from "Shared/Telemetry/TelemetryConstants";
+import { traceFailure, traceStart, traceSuccess } from "Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "UserContext";
 import { queryPagesUntilContentPresent } from "Utils/QueryUtils";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
@@ -78,6 +80,7 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({
   const hideFeedbackModalForLikedQueries = useQueryCopilot((state) => state.hideFeedbackModalForLikedQueries);
   const [userPrompt, setUserPrompt] = useState<string>(initialInput || "");
   const [generatedQuery, setGeneratedQuery] = useState<string>("");
+  const [generatedQueryComments, setGeneratedQueryComments] = useState<string>("");
   const [query, setQuery] = useState<string>("");
   const [selectedQuery, setSelectedQuery] = useState<string>("");
   const [isGeneratingQuery, setIsGeneratingQuery] = useState<boolean>(false);
@@ -170,10 +173,12 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({
         userPrompt: userPrompt,
       };
       setShowDeletePopup(false);
+      useQueryCopilot.getState().refreshCorrelationId();
       const response = await fetch("https://copilotorchestrater.azurewebsites.net/generateSQLQuery", {
         method: "POST",
         headers: {
           "content-type": "application/json",
+          "x-ms-correlationid": useQueryCopilot.getState().correlationId,
         },
         body: JSON.stringify(payload),
       });
@@ -188,6 +193,7 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({
           query += generateSQLQueryResponse.sql;
           setQuery(query);
           setGeneratedQuery(generateSQLQueryResponse.sql);
+          setGeneratedQueryComments(generateSQLQueryResponse.explanation);
           setShowErrorMessageBar(false);
         }
       } else {
@@ -208,6 +214,13 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({
   };
 
   const onExecuteQueryClick = async (): Promise<void> => {
+    traceStart(Action.ExecuteQueryGeneratedFromQueryCopilot, {
+      correlationId: useQueryCopilot.getState().correlationId,
+      userPrompt: userPrompt,
+      generatedQuery: generatedQuery,
+      generatedQueryComments: generatedQueryComments,
+      executedQuery: selectedQuery || query,
+    });
     const queryToExecute = selectedQuery || query;
     const queryIterator = querySampleDocuments(queryToExecute, {
       enableCrossPartitionQuery: shouldEnableCrossPartitionKey(),
@@ -233,8 +246,15 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({
       setQueryResults(queryResults);
       setErrorMessage("");
       setShowErrorMessageBar(false);
+      traceSuccess(Action.ExecuteQueryGeneratedFromQueryCopilot, {
+        correlationId: useQueryCopilot.getState().correlationId,
+      });
     } catch (error) {
       const errorMessage = getErrorMessage(error);
+      traceFailure(Action.ExecuteQueryGeneratedFromQueryCopilot, {
+        correlationId: useQueryCopilot.getState().correlationId,
+        errorMessage: errorMessage,
+      });
       setErrorMessage(errorMessage);
       handleError(errorMessage, "executeQueryCopilotTab");
       useTabs.getState().setIsQueryErrorThrown(true);

src/Explorer/QueryCopilot/QueryCopilotUtilities.test.tsx

@@ -4,6 +4,7 @@ import { handleError } from "Common/ErrorHandlingUtils";
 import { sampleDataClient } from "Common/SampleDataClient";
 import * as commonUtils from "Common/dataAccess/queryDocuments";
 import DocumentId from "Explorer/Tree/DocumentId";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { querySampleDocuments, readSampleDocument, submitFeedback } from "./QueryCopilotUtilities";
 jest.mock("Explorer/Tree/DocumentId", () => {
   return jest.fn().mockImplementation(() => {
@@ -56,6 +57,7 @@ describe("QueryCopilotUtilities", () => {
       const mockFetch = jest.fn().mockResolvedValueOnce({});
 
       globalThis.fetch = mockFetch;
+      useQueryCopilot.getState().refreshCorrelationId();
 
       await submitFeedback({
         likeQuery: true,
@@ -71,6 +73,7 @@ describe("QueryCopilotUtilities", () => {
           method: "POST",
           headers: {
             "content-type": "application/json",
+            "x-ms-correlationid": useQueryCopilot.getState().correlationId,
           },
         })
       );
@@ -86,6 +89,7 @@ describe("QueryCopilotUtilities", () => {
       const mockFetch = jest.fn().mockResolvedValueOnce({});
 
       globalThis.fetch = mockFetch;
+      useQueryCopilot.getState().refreshCorrelationId();
 
       await submitFeedback({
         likeQuery: false,
@@ -101,6 +105,7 @@ describe("QueryCopilotUtilities", () => {
           method: "POST",
           headers: {
             "content-type": "application/json",
+            "x-ms-correlationid": useQueryCopilot.getState().correlationId,
           },
         })
       );

src/Explorer/QueryCopilot/QueryCopilotUtilities.ts

@@ -10,6 +10,7 @@ import { getPartitionKeyValue } from "Common/dataAccess/getPartitionKeyValue";
 import { getCommonQueryOptions } from "Common/dataAccess/queryDocuments";
 import DocumentId from "Explorer/Tree/DocumentId";
 import { logConsoleProgress } from "Utils/NotificationConsoleUtils";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 
 interface FeedbackParams {
   likeQuery: boolean;
@@ -35,6 +36,7 @@ export const submitFeedback = async (params: FeedbackParams): Promise<void> => {
       method: "POST",
       headers: {
         "content-type": "application/json",
+        "x-ms-correlationid": useQueryCopilot.getState().correlationId,
       },
       body: JSON.stringify(payload),
     });

src/Explorer/SplashScreen/SplashScreen.tsx

@@ -137,7 +137,10 @@ export class SplashScreen extends React.Component<SplashScreenProps> {
               description={
                 "Copilot is your AI buddy that helps you write Azure Cosmos DB queries like a pro. Try it using our sample data set now!"
               }
-              onClick={() => useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot)}
+              onClick={() => {
+                useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot);
+                traceOpen(Action.OpenQueryCopilotFromSplashScreen, { apiType: userContext.apiType });
+              }}
             />
             <SplashScreenButton
               imgSrc={ConnectIcon}

src/Shared/Telemetry/TelemetryConstants.ts

@@ -131,6 +131,9 @@ export enum Action {
   LaunchUITour,
   CancelUITour,
   CompleteUITour,
+  OpenQueryCopilotFromSplashScreen,
+  OpenQueryCopilotFromNewQuery,
+  ExecuteQueryGeneratedFromQueryCopilot,
 }
 
 export const ActionModifiers = {

src/hooks/useQueryCopilot.ts

@@ -1,3 +1,4 @@
+import { guid } from "Explorer/Tables/Utilities";
 import create, { UseStore } from "zustand";
 
 interface QueryCopilotState {
@@ -6,9 +7,11 @@ interface QueryCopilotState {
   userPrompt: string;
   showFeedbackModal: boolean;
   hideFeedbackModalForLikedQueries: boolean;
+  correlationId: string;
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void;
   closeFeedbackModal: () => void;
   setHideFeedbackModalForLikedQueries: (hideFeedbackModalForLikedQueries: boolean) => void;
+  refreshCorrelationId: () => void;
 }
 
 export const useQueryCopilot: UseStore<QueryCopilotState> = create((set) => ({
@@ -17,9 +20,11 @@ export const useQueryCopilot: UseStore<QueryCopilotState> = create((set) => ({
   userPrompt: "",
   showFeedbackModal: false,
   hideFeedbackModalForLikedQueries: false,
+  correlationId: "",
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
     set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
   closeFeedbackModal: () => set({ generatedQuery: "", likeQuery: false, userPrompt: "", showFeedbackModal: false }),
   setHideFeedbackModalForLikedQueries: (hideFeedbackModalForLikedQueries: boolean) =>
     set({ hideFeedbackModalForLikedQueries }),
+  refreshCorrelationId: () => set({ correlationId: guid() }),
 }));