[Query Copilot] Allocation of container and copilot request sent to Phoenix endpoint (#1576)

* Switch to tools federation endpoints for query copilot

* Remove extra / in url

* Initial allocateContainer implementation

* Additional feedback modal changes

* updated tests

* PhoenixClient reverted to previous endpoint

* Changes based on PR comments

* Update based on tests

* updated endpoint

* Back to previous implementation and test improve

* removing notebook

---------

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

src/Common/Constants.ts

@@ -358,6 +358,7 @@ export enum ContainerStatusType {
 
 export enum PoolIdType {
   DefaultPoolId = "default",
+  QueryCopilot = "query-copilot",
 }
 
 export const EmulatorMasterKey =

src/Explorer/Explorer.tsx

@@ -395,7 +395,7 @@ export default class Explorer {
     ) {
       const provisionData: IProvisionData = {
         cosmosEndpoint: userContext?.databaseAccount?.properties?.documentEndpoint,
-        poolId: PoolIdType.DefaultPoolId,
+        poolId: PoolIdType.QueryCopilot,
       };
       const connectionStatus: ContainerConnectionInfo = {
         status: ConnectionStatusType.Connecting,

src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx

@@ -1,4 +1,5 @@
 import { Checkbox, ChoiceGroup, DefaultButton, IconButton, PrimaryButton, TextField } from "@fluentui/react";
+import Explorer from "Explorer/Explorer";
 import { QueryCopilotFeedbackModal } from "Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal";
 import { submitFeedback } from "Explorer/QueryCopilot/QueryCopilotUtilities";
 import { getUserEmail } from "Utils/UserUtils";
@@ -19,14 +20,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   it("shoud render and match snapshot", () => {
     useQueryCopilot.getState().openFeedbackModal("test query", false, "test prompt");
 
-    const wrapper = shallow(<QueryCopilotFeedbackModal />);
+    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
 
     expect(wrapper.props().isOpen).toBeTruthy();
     expect(wrapper).toMatchSnapshot();
   });
 
   it("should close on cancel click", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal />);
+    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
 
     const cancelButton = wrapper.find(IconButton);
     cancelButton.simulate("click");
@@ -37,7 +38,7 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should get user unput", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal />);
+    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
     const testUserInput = "test user input";
 
     const userInput = wrapper.find(TextField).first();
@@ -48,7 +49,7 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should record user contact choice no", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal />);
+    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
     const contactAllowed = wrapper.find(ChoiceGroup);
 
     contactAllowed.simulate("change", {}, { key: "no" });
@@ -59,7 +60,7 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should record user contact choice yes", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal />);
+    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
     const contactAllowed = wrapper.find(ChoiceGroup);
 
     contactAllowed.simulate("change", {}, { key: "yes" });
@@ -70,7 +71,7 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should not render dont show again button", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal />);
+    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
 
     const dontShowAgain = wrapper.find(Checkbox);
 
@@ -80,7 +81,7 @@ 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(<QueryCopilotFeedbackModal />);
+    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
 
     const dontShowAgain = wrapper.find(Checkbox);
     dontShowAgain.simulate("change", {}, true);
@@ -91,7 +92,7 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should cancel submission", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal />);
+    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
 
     const cancelButton = wrapper.find(DefaultButton);
     cancelButton.simulate("click");
@@ -102,7 +103,8 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should submit submission", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal />);
+    const explorer = new Explorer();
+    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={explorer} />);
 
     const submitButton = wrapper.find(PrimaryButton);
     submitButton.simulate("click");
@@ -110,11 +112,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
 
     expect(submitFeedback).toHaveBeenCalledTimes(1);
     expect(submitFeedback).toHaveBeenCalledWith({
-      likeQuery: false,
+      params: {
-      generatedQuery: "",
+        likeQuery: false,
-      userPrompt: "",
+        generatedQuery: "",
-      description: "",
+        userPrompt: "",
-      contact: getUserEmail(),
+        description: "",
+        contact: getUserEmail(),
+      },
+      explorer: explorer,
     });
     expect(wrapper.props().isOpen).toBeFalsy();
     expect(wrapper).toMatchSnapshot();

src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx

@@ -10,12 +10,13 @@ import {
   Text,
   TextField,
 } from "@fluentui/react";
+import Explorer from "Explorer/Explorer";
 import { submitFeedback } from "Explorer/QueryCopilot/QueryCopilotUtilities";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
 import React from "react";
 import { getUserEmail } from "../../../Utils/UserUtils";
 
-export const QueryCopilotFeedbackModal: React.FC = (): JSX.Element => {
+export const QueryCopilotFeedbackModal = ({ explorer }: { explorer: Explorer }): JSX.Element => {
   const {
     generatedQuery,
     userPrompt,
@@ -100,7 +101,10 @@ export const QueryCopilotFeedbackModal: React.FC = (): JSX.Element => {
             onClick={() => {
               closeFeedbackModal();
               setHideFeedbackModalForLikedQueries(doNotShowAgainChecked);
-              submitFeedback({ generatedQuery, likeQuery, description, userPrompt, contact });
+              submitFeedback({
+                params: { generatedQuery, likeQuery, description, userPrompt, contact },
+                explorer: explorer,
+              });
             }}
           >
             Submit

src/Explorer/QueryCopilot/QueryCopilotTab.tsx

@@ -21,12 +21,14 @@ import { QueryCopilotSampleContainerId, QueryCopilotSampleContainerSchema } from
 import { getErrorMessage, 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 { CommandButtonComponentProps } from "Explorer/Controls/CommandButton/CommandButtonComponent";
 import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
 import Explorer from "Explorer/Explorer";
 import { useCommandBar } from "Explorer/Menus/CommandBar/CommandBarComponentAdapter";
+import { useNotebook } from "Explorer/Notebook/useNotebook";
 import { SaveQueryPane } from "Explorer/Panes/SaveQueryPane/SaveQueryPane";
 import { WelcomeModal } from "Explorer/QueryCopilot/Modal/WelcomeModal";
 import { CopyPopup } from "Explorer/QueryCopilot/Popup/CopyPopup";
@@ -115,6 +117,8 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({ explorer }: Qu
     setShowErrorMessageBar,
     generatedQueryComments,
     setGeneratedQueryComments,
+    shouldAllocateContainer,
+    setShouldAllocateContainer,
   } = useQueryCopilot();
 
   const sampleProps: SamplePromptsProps = {
@@ -182,6 +186,11 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({ explorer }: Qu
 
   const generateSQLQuery = async (): Promise<void> => {
     try {
+      if (shouldAllocateContainer) {
+        await explorer.allocateContainer();
+        setShouldAllocateContainer(false);
+      }
+
       setIsGeneratingQuery(true);
       useTabs.getState().setIsTabExecuting(true);
       useTabs.getState().setIsQueryErrorThrown(false);
@@ -191,7 +200,9 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({ explorer }: Qu
       };
       setShowDeletePopup(false);
       useQueryCopilot.getState().refreshCorrelationId();
-      const response = await fetch("https://copilotorchestrater.azurewebsites.net/generateSQLQuery", {
+      const serverInfo = useNotebook.getState().notebookServerInfo;
+      const queryUri = createUri(serverInfo.notebookServerEndpoint, "generateSQLQuery");
+      const response = await fetch(queryUri, {
         method: "POST",
         headers: {
           "content-type": "application/json",
@@ -201,6 +212,9 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({ explorer }: Qu
       });
 
       const generateSQLQueryResponse: GenerateSQLQueryResponse = await response?.json();
+      if (response.status === 404) {
+        setShouldAllocateContainer(true);
+      }
       if (response.ok) {
         if (generateSQLQueryResponse?.sql) {
           let query = `-- **Prompt:** ${userPrompt}\r\n`;
@@ -533,10 +547,13 @@ export const QueryCopilotTab: React.FC<QueryCopilotTabProps> = ({ explorer }: Qu
                 onDismiss={() => {
                   setShowCallout(false);
                   submitFeedback({
-                    generatedQuery: generatedQuery,
+                    params: {
-                    likeQuery: likeQuery,
+                      generatedQuery: generatedQuery,
-                    description: "",
+                      likeQuery: likeQuery,
-                    userPrompt: userPrompt,
+                      description: "",
+                      userPrompt: userPrompt,
+                    },
+                    explorer: explorer,
                   });
                 }}
                 directionalHint={DirectionalHint.topCenter}

src/Explorer/QueryCopilot/QueryCopilotUtilities.test.tsx

@@ -3,9 +3,11 @@ import { QueryCopilotSampleContainerSchema } from "Common/Constants";
 import { handleError } from "Common/ErrorHandlingUtils";
 import { sampleDataClient } from "Common/SampleDataClient";
 import * as commonUtils from "Common/dataAccess/queryDocuments";
+import Explorer from "Explorer/Explorer";
+import { useNotebook } from "Explorer/Notebook/useNotebook";
 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(() => {
     return {
@@ -16,23 +18,21 @@ jest.mock("Explorer/Tree/DocumentId", () => {
 });
 
 jest.mock("Utils/NotificationConsoleUtils", () => ({
-  logConsoleProgress: jest.fn(),
+  logConsoleProgress: jest.fn().mockReturnValue((): void => undefined),
   logConsoleError: jest.fn(),
 }));
 
 jest.mock("@azure/cosmos", () => ({
   FeedOptions: jest.fn(),
   QueryIterator: jest.fn(),
+  Constants: {
+    HttpHeaders: {},
+  },
 }));
-
 jest.mock("Common/ErrorHandlingUtils", () => ({
   handleError: jest.fn(),
 }));
 
-jest.mock("Utils/NotificationConsoleUtils", () => ({
-  logConsoleProgress: jest.fn().mockReturnValue((): void => undefined),
-}));
-
 jest.mock("Common/dataAccess/queryDocuments", () => ({
   getCommonQueryOptions: jest.fn((options) => options),
 }));
@@ -41,8 +41,28 @@ jest.mock("Common/SampleDataClient");
 
 jest.mock("node-fetch");
 
+jest.mock("Explorer/Explorer", () => {
+  class MockExplorer {
+    allocateContainer = jest.fn().mockResolvedValueOnce({});
+  }
+  return MockExplorer;
+});
+
+jest.mock("hooks/useQueryCopilot", () => {
+  const mockQueryCopilotStore = {
+    shouldAllocateContainer: true,
+    setShouldAllocateContainer: jest.fn(),
+    correlationId: "mocked-correlation-id",
+  };
+
+  return {
+    useQueryCopilot: jest.fn(() => mockQueryCopilotStore),
+  };
+});
+
 describe("QueryCopilotUtilities", () => {
   beforeEach(() => jest.clearAllMocks());
+
   describe("submitFeedback", () => {
     const payload = {
       like: "like",
@@ -53,28 +73,37 @@ describe("QueryCopilotUtilities", () => {
       containerSchema: QueryCopilotSampleContainerSchema,
     };
 
+    const mockStore = useNotebook.getState();
+    beforeEach(() => {
+      mockStore.notebookServerInfo = {
+        notebookServerEndpoint: "mocked-endpoint",
+        authToken: "mocked-token",
+        forwardingId: "mocked-forwarding-id",
+      };
+    });
+
     it("should call fetch with the payload with like", async () => {
       const mockFetch = jest.fn().mockResolvedValueOnce({});
 
       globalThis.fetch = mockFetch;
-      useQueryCopilot.getState().refreshCorrelationId();
 
       await submitFeedback({
-        likeQuery: true,
+        params: {
-        generatedQuery: "GeneratedQuery",
+          likeQuery: true,
-        userPrompt: "UserPrompt",
+          generatedQuery: "GeneratedQuery",
-        description: "Description",
+          userPrompt: "UserPrompt",
-        contact: "Contact",
+          description: "Description",
+          contact: "Contact",
+        },
+        explorer: new Explorer(),
       });
 
       expect(mockFetch).toHaveBeenCalledWith(
-        "https://copilotorchestrater.azurewebsites.net/feedback",
+        "mocked-endpoint/feedback",
         expect.objectContaining({
-          method: "POST",
+          headers: expect.objectContaining({
-          headers: {
+            "x-ms-correlationid": "mocked-correlation-id",
-            "content-type": "application/json",
+          }),
-            "x-ms-correlationid": useQueryCopilot.getState().correlationId,
-          },
         })
       );
 
@@ -89,23 +118,25 @@ describe("QueryCopilotUtilities", () => {
       const mockFetch = jest.fn().mockResolvedValueOnce({});
 
       globalThis.fetch = mockFetch;
-      useQueryCopilot.getState().refreshCorrelationId();
 
       await submitFeedback({
-        likeQuery: false,
+        params: {
-        generatedQuery: "GeneratedQuery",
+          likeQuery: false,
-        userPrompt: "UserPrompt",
+          generatedQuery: "GeneratedQuery",
-        description: undefined,
+          userPrompt: "UserPrompt",
-        contact: undefined,
+          description: undefined,
+          contact: undefined,
+        },
+        explorer: new Explorer(),
       });
 
       expect(mockFetch).toHaveBeenCalledWith(
-        "https://copilotorchestrater.azurewebsites.net/feedback",
+        "mocked-endpoint/feedback",
         expect.objectContaining({
           method: "POST",
           headers: {
             "content-type": "application/json",
-            "x-ms-correlationid": useQueryCopilot.getState().correlationId,
+            "x-ms-correlationid": "mocked-correlation-id",
           },
         })
       );
@@ -118,11 +149,14 @@ describe("QueryCopilotUtilities", () => {
       globalThis.fetch = jest.fn().mockRejectedValueOnce(new Error("Mock error"));
 
       await submitFeedback({
-        likeQuery: true,
+        params: {
-        generatedQuery: "GeneratedQuery",
+          likeQuery: true,
-        userPrompt: "UserPrompt",
+          generatedQuery: "GeneratedQuery",
-        description: "Description",
+          userPrompt: "UserPrompt",
-        contact: "Contact",
+          description: "Description",
+          contact: "Contact",
+        },
+        explorer: new Explorer(),
       }).catch((error) => {
         expect(error.message).toEqual("Mock error");
       });

src/Explorer/QueryCopilot/QueryCopilotUtilities.ts

@@ -6,8 +6,11 @@ import {
 } from "Common/Constants";
 import { handleError } from "Common/ErrorHandlingUtils";
 import { sampleDataClient } from "Common/SampleDataClient";
+import { createUri } from "Common/UrlUtility";
 import { getPartitionKeyValue } from "Common/dataAccess/getPartitionKeyValue";
 import { getCommonQueryOptions } from "Common/dataAccess/queryDocuments";
+import Explorer from "Explorer/Explorer";
+import { useNotebook } from "Explorer/Notebook/useNotebook";
 import DocumentId from "Explorer/Tree/DocumentId";
 import { logConsoleProgress } from "Utils/NotificationConsoleUtils";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
@@ -20,9 +23,16 @@ interface FeedbackParams {
   contact?: string;
 }
 
-export const submitFeedback = async (params: FeedbackParams): Promise<void> => {
+export const submitFeedback = async ({
+  params,
+  explorer,
+}: {
+  params: FeedbackParams;
+  explorer: Explorer;
+}): Promise<void> => {
   try {
     const { likeQuery, generatedQuery, userPrompt, description, contact } = params;
+    const { correlationId, shouldAllocateContainer, setShouldAllocateContainer } = useQueryCopilot();
     const payload = {
       containerSchema: QueryCopilotSampleContainerSchema,
       like: likeQuery ? "like" : "dislike",
@@ -31,15 +41,23 @@ export const submitFeedback = async (params: FeedbackParams): Promise<void> => {
       description: description || "",
       contact: contact || "",
     };
-
+    if (shouldAllocateContainer) {
-    await fetch("https://copilotorchestrater.azurewebsites.net/feedback", {
+      await explorer.allocateContainer();
+      setShouldAllocateContainer(false);
+    }
+    const serverInfo = useNotebook.getState().notebookServerInfo;
+    const feedbackUri = createUri(serverInfo.notebookServerEndpoint, "feedback");
+    const response = await fetch(feedbackUri, {
       method: "POST",
       headers: {
         "content-type": "application/json",
-        "x-ms-correlationid": useQueryCopilot.getState().correlationId,
+        "x-ms-correlationid": correlationId,
       },
       body: JSON.stringify(payload),
     });
+    if (response.status === 404) {
+      setShouldAllocateContainer(true);
+    }
   } catch (error) {
     handleError(error, "copilotSubmitFeedback");
   }

src/Main.tsx

@@ -128,7 +128,7 @@ const App: React.FunctionComponent = () => {
       {<SQLQuickstartTutorial />}
       {<MongoQuickstartTutorial />}
       {<QueryCopilotCarousel isOpen={isCopilotCarouselOpen} explorer={explorer} />}
-      {shouldShowModal && <QueryCopilotFeedbackModal />}
+      {shouldShowModal && <QueryCopilotFeedbackModal explorer={explorer} />}
     </div>
   );
 };

src/Phoenix/PhoenixClient.ts

@@ -1,7 +1,9 @@
+import { configContext } from "ConfigContext";
 import { useDialog } from "Explorer/Controls/Dialog";
-import promiseRetry, { AbortError } from "p-retry";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
+import { userContext } from "UserContext";
 import { allowedJunoOrigins, validateEndpoint } from "Utils/EndpointValidation";
+import promiseRetry, { AbortError } from "p-retry";
 import {
   Areas,
   ConnectionStatusType,
@@ -12,7 +14,6 @@ import {
 } from "../Common/Constants";
 import { getErrorMessage, getErrorStack } from "../Common/ErrorHandlingUtils";
 import * as Logger from "../Common/Logger";
-import { configContext } from "../ConfigContext";
 import {
   ContainerConnectionInfo,
   ContainerInfo,
@@ -28,7 +29,6 @@ import {
 } from "../Contracts/DataModels";
 import { useNotebook } from "../Explorer/Notebook/useNotebook";
 import * as TelemetryProcessor from "../Shared/Telemetry/TelemetryProcessor";
-import { userContext } from "../UserContext";
 import { getAuthorizationHeader } from "../Utils/AuthorizationUtils";
 
 export class PhoenixClient {

src/hooks/useQueryCopilot.ts

@@ -30,6 +30,7 @@ export interface QueryCopilotState {
   showWelcomeSidebar: boolean;
   showCopilotSidebar: boolean;
   chatMessages: string[];
+  shouldAllocateContainer: boolean;
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void;
   closeFeedbackModal: () => void;
@@ -59,6 +60,8 @@ export interface QueryCopilotState {
   setShowCopilotSidebar: (showCopilotSidebar: boolean) => void;
   setChatMessages: (chatMessages: string[]) => void;
 
+  setShouldAllocateContainer: (shouldAllocateContainer: boolean) => void;
+
   resetQueryCopilotStates: () => void;
 }
 
@@ -91,6 +94,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   showWelcomeSidebar: true,
   showCopilotSidebar: false,
   chatMessages: [],
+  shouldAllocateContainer: true,
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
     set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
@@ -121,6 +125,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => set({ showWelcomeSidebar }),
   setShowCopilotSidebar: (showCopilotSidebar: boolean) => set({ showCopilotSidebar }),
   setChatMessages: (chatMessages: string[]) => set({ chatMessages }),
+  setShouldAllocateContainer: (shouldAllocateContainer: boolean) => set({ shouldAllocateContainer }),
 
   resetQueryCopilotStates: () => {
     set((state) => ({
@@ -150,6 +155,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       wasCopilotUsed: false,
       showCopilotSidebar: false,
       chatMessages: [],
+      shouldAllocateContainer: true,
     }));
   },
 }));