random-mirrors
/
cosmos-explorer
mirror of https://github.com/Azure/cosmos-explorer.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[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>
This commit is contained in:
Predrag Klepic 2023-08-18 10:47:19 +02:00 committed by GitHub
parent 19041ffedd
commit 9c1b9e6ff6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 147 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/Common/Constants.ts
View File

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

2
src/Explorer/Explorer.tsx
View File

@ -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,

33
src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx
View File

@ -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,
generatedQuery: "",
userPrompt: "",
description: "",
contact: getUserEmail(),
params: {
likeQuery: false,
generatedQuery: "",
userPrompt: "",
description: "",
contact: getUserEmail(),
},
explorer: explorer,
});
expect(wrapper.props().isOpen).toBeFalsy();
expect(wrapper).toMatchSnapshot();

8
src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx
View File

@ -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

27
src/Explorer/QueryCopilot/QueryCopilotTab.tsx
View File

@ -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,
likeQuery: likeQuery,
description: "",
userPrompt: userPrompt,
params: {
generatedQuery: generatedQuery,
likeQuery: likeQuery,
description: "",
userPrompt: userPrompt,
},
explorer: explorer,
});
}}
directionalHint={DirectionalHint.topCenter}

98
src/Explorer/QueryCopilot/QueryCopilotUtilities.test.tsx
View File

@ -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");
});

26
src/Explorer/QueryCopilot/QueryCopilotUtilities.ts
View File

@ -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 || "",
};
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");
}

2
src/Main.tsx
View File

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

6
src/Phoenix/PhoenixClient.ts
View File

@ -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 {

6
src/hooks/useQueryCopilot.ts
View File

@ -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,
}));
},
}));