diff --git a/src/Common/Constants.ts b/src/Common/Constants.ts index 26d0d8556..352378a70 100644 --- a/src/Common/Constants.ts +++ b/src/Common/Constants.ts @@ -358,6 +358,7 @@ export enum ContainerStatusType { export enum PoolIdType { DefaultPoolId = "default", + QueryCopilot = "query-copilot", } export const EmulatorMasterKey = diff --git a/src/Explorer/Explorer.tsx b/src/Explorer/Explorer.tsx index 824be1ee2..59d8847ef 100644 --- a/src/Explorer/Explorer.tsx +++ b/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, diff --git a/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx b/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx index 17ebfa422..096692105 100644 --- a/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx +++ b/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(); + 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"); @@ -37,7 +38,7 @@ 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(); @@ -48,7 +49,7 @@ 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" }); @@ -59,7 +60,7 @@ 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" }); @@ -70,7 +71,7 @@ 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,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(); + const wrapper = shallow(); 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(); + const wrapper = shallow(); 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(); + const explorer = new Explorer(); + const wrapper = shallow(); 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, - generatedQuery: "", - userPrompt: "", - description: "", - contact: getUserEmail(), + params: { + likeQuery: false, + generatedQuery: "", + userPrompt: "", + description: "", + contact: getUserEmail(), + }, + explorer: explorer, }); expect(wrapper.props().isOpen).toBeFalsy(); expect(wrapper).toMatchSnapshot(); diff --git a/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx b/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx index ac63b178b..d3eaa47b6 100644 --- a/src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx +++ b/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 diff --git a/src/Explorer/QueryCopilot/QueryCopilotTab.tsx b/src/Explorer/QueryCopilot/QueryCopilotTab.tsx index 0ae18f751..1c32596b4 100644 --- a/src/Explorer/QueryCopilot/QueryCopilotTab.tsx +++ b/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 = ({ explorer }: Qu setShowErrorMessageBar, generatedQueryComments, setGeneratedQueryComments, + shouldAllocateContainer, + setShouldAllocateContainer, } = useQueryCopilot(); const sampleProps: SamplePromptsProps = { @@ -182,6 +186,11 @@ export const QueryCopilotTab: React.FC = ({ explorer }: Qu const generateSQLQuery = async (): Promise => { 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 = ({ 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 = ({ 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 = ({ explorer }: Qu onDismiss={() => { setShowCallout(false); submitFeedback({ - generatedQuery: generatedQuery, - likeQuery: likeQuery, - description: "", - userPrompt: userPrompt, + params: { + generatedQuery: generatedQuery, + likeQuery: likeQuery, + description: "", + userPrompt: userPrompt, + }, + explorer: explorer, }); }} directionalHint={DirectionalHint.topCenter} diff --git a/src/Explorer/QueryCopilot/QueryCopilotUtilities.test.tsx b/src/Explorer/QueryCopilot/QueryCopilotUtilities.test.tsx index ee09d1820..1be2873cb 100644 --- a/src/Explorer/QueryCopilot/QueryCopilotUtilities.test.tsx +++ b/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, - generatedQuery: "GeneratedQuery", - userPrompt: "UserPrompt", - description: "Description", - contact: "Contact", + params: { + likeQuery: true, + generatedQuery: "GeneratedQuery", + userPrompt: "UserPrompt", + description: "Description", + contact: "Contact", + }, + 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, - }, + headers: expect.objectContaining({ + "x-ms-correlationid": "mocked-correlation-id", + }), }) ); @@ -89,23 +118,25 @@ describe("QueryCopilotUtilities", () => { const mockFetch = jest.fn().mockResolvedValueOnce({}); globalThis.fetch = mockFetch; - useQueryCopilot.getState().refreshCorrelationId(); await submitFeedback({ - likeQuery: false, - generatedQuery: "GeneratedQuery", - userPrompt: "UserPrompt", - description: undefined, - contact: undefined, + params: { + likeQuery: false, + generatedQuery: "GeneratedQuery", + userPrompt: "UserPrompt", + 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, - generatedQuery: "GeneratedQuery", - userPrompt: "UserPrompt", - description: "Description", - contact: "Contact", + params: { + likeQuery: true, + generatedQuery: "GeneratedQuery", + userPrompt: "UserPrompt", + description: "Description", + contact: "Contact", + }, + explorer: new Explorer(), }).catch((error) => { expect(error.message).toEqual("Mock error"); }); diff --git a/src/Explorer/QueryCopilot/QueryCopilotUtilities.ts b/src/Explorer/QueryCopilot/QueryCopilotUtilities.ts index f4af4fb97..9a3c5a33d 100644 --- a/src/Explorer/QueryCopilot/QueryCopilotUtilities.ts +++ b/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 => { +export const submitFeedback = async ({ + params, + explorer, +}: { + params: FeedbackParams; + explorer: Explorer; +}): Promise => { 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 => { description: description || "", contact: contact || "", }; - - await fetch("https://copilotorchestrater.azurewebsites.net/feedback", { + if (shouldAllocateContainer) { + 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"); } diff --git a/src/Main.tsx b/src/Main.tsx index a7f35eb6a..6ea26ba9f 100644 --- a/src/Main.tsx +++ b/src/Main.tsx @@ -128,7 +128,7 @@ const App: React.FunctionComponent = () => { {} {} {} - {shouldShowModal && } + {shouldShowModal && } ); }; diff --git a/src/Phoenix/PhoenixClient.ts b/src/Phoenix/PhoenixClient.ts index ddc4d563a..ff71e50e3 100644 --- a/src/Phoenix/PhoenixClient.ts +++ b/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 { diff --git a/src/hooks/useQueryCopilot.ts b/src/hooks/useQueryCopilot.ts index 3ead839b6..7965592cd 100644 --- a/src/hooks/useQueryCopilot.ts +++ b/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, })); }, }));