random-mirrors/cosmos-explorer
1
0
Fork 0
mirror of https://github.com/Azure/cosmos-explorer.git synced 2025-12-23 02:41:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Query Copilot] Phoenix container implementation (#1594)

Browse Source 
* Phoenix implementation

* removing comments

---------

Co-authored-by: Predrag Klepic <v-prklepic@microsoft.com>
This commit is contained in:
Predrag Klepic
2023-08-30 16:50:55 +02:00
committed by GitHub
parent b992742e20
commit c8e7e69aa5
6 changed files with 158 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -18,6 +18,7 @@ import {
} from "@fluentui/react";
import { useBoolean } from "@fluentui/react-hooks";
import {
ContainerStatusType,
PoolIdType,
QueryCopilotSampleContainerId,
QueryCopilotSampleContainerSchema,
@@ -32,7 +33,6 @@ import { QueryResults } from "Contracts/ViewModels";
import { CommandButtonComponentProps } from "Explorer/Controls/CommandButton/CommandButtonComponent";
import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
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";
@@ -111,8 +111,6 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
setShowErrorMessageBar,
generatedQueryComments,
setGeneratedQueryComments,
shouldAllocateContainer,
setShouldAllocateContainer,
} = useQueryCopilot();
const sampleProps: SamplePromptsProps = {
@@ -180,23 +178,24 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
const generateSQLQuery = async (): Promise<void> => {
try {
if (shouldAllocateContainer && userContext.features.enableCopilotPhoenixGateaway) {
await explorer.allocateContainer(PoolIdType.QueryCopilot);
setShouldAllocateContainer(false);
}
setIsGeneratingQuery(true);
setShowDeletePopup(false);
useTabs.getState().setIsTabExecuting(true);
useTabs.getState().setIsQueryErrorThrown(false);
if (
useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
userContext.features.enableCopilotPhoenixGateaway
) {
await explorer.allocateContainer(PoolIdType.QueryCopilot);
}
const payload = {
containerSchema: userContext.features.enableCopilotFullSchema
? QueryCopilotSampleContainerSchema
: ShortenedQueryCopilotSampleContainerSchema,
userPrompt: userPrompt,
};
setShowDeletePopup(false);
useQueryCopilot.getState().refreshCorrelationId();
const serverInfo = useNotebook.getState().notebookServerInfo;
const serverInfo = useQueryCopilot.getState().notebookServerInfo;
const queryUri = userContext.features.enableCopilotPhoenixGateaway
? createUri(serverInfo.notebookServerEndpoint, "generateSQLQuery")
: createUri("https://copilotorchestrater.azurewebsites.net/", "generateSQLQuery");
@@ -210,9 +209,6 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
});
const generateSQLQueryResponse: GenerateSQLQueryResponse = await response?.json();
if (response.status === 404) {
setShouldAllocateContainer(true);
}
if (response.ok) {
if (generateSQLQueryResponse?.sql) {
let query = `-- **Prompt:** ${userPrompt}\r\n`;

33
src/Explorer/QueryCopilot/Shared/QueryCopilotClient.test.ts
View File

@@ -2,9 +2,9 @@ import { QueryCopilotSampleContainerSchema, ShortenedQueryCopilotSampleContainer
import { handleError } from "Common/ErrorHandlingUtils";
import { createUri } from "Common/UrlUtility";
import Explorer from "Explorer/Explorer";
import { useNotebook } from "Explorer/Notebook/useNotebook";
import { SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
import { userContext } from "UserContext";
import { useQueryCopilot } from "hooks/useQueryCopilot";
jest.mock("@azure/cosmos", () => ({
Constants: {
@@ -27,22 +27,10 @@ jest.mock("Explorer/Explorer", () => {
return MockExplorer;
});
jest.mock("hooks/useQueryCopilot", () => {
const mockQueryCopilotStore = {
shouldAllocateContainer: true,
setShouldAllocateContainer: jest.fn(),
correlationId: "mocked-correlation-id",
};
return {
useQueryCopilot: jest.fn(() => mockQueryCopilotStore),
};
});
describe("Query Copilot Client", () => {
beforeEach(() => jest.clearAllMocks());
describe("submitFeedback", () => {
describe("SubmitFeedback", () => {
const payload = {
like: "like",
generatedSql: "GeneratedQuery",
@@ -54,17 +42,16 @@ describe("Query Copilot Client", () => {
: ShortenedQueryCopilotSampleContainerSchema,
};
const mockStore = useNotebook.getState();
beforeEach(() => {
mockStore.notebookServerInfo = {
notebookServerEndpoint: "mocked-endpoint",
authToken: "mocked-token",
forwardingId: "mocked-forwarding-id",
};
});
const mockStore = useQueryCopilot.getState();
mockStore.correlationId = "mocked-correlation-id";
mockStore.notebookServerInfo = {
notebookServerEndpoint: "mocked-endpoint",
authToken: "mocked-token",
forwardingId: "mocked-forwarding-id",
};
const feedbackUri = userContext.features.enableCopilotPhoenixGateaway
? createUri(useNotebook.getState().notebookServerInfo.notebookServerEndpoint, "feedback")
? createUri(useQueryCopilot.getState().notebookServerInfo.notebookServerEndpoint, "feedback")
: createUri("https://copilotorchestrater.azurewebsites.net/", "feedback");
it("should call fetch with the payload with like", async () => {

33
src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts
View File

@@ -1,4 +1,5 @@
import {
ContainerStatusType,
PoolIdType,
QueryCopilotSampleContainerSchema,
ShortenedQueryCopilotSampleContainerSchema,
@@ -6,7 +7,6 @@ import {
import { handleError } from "Common/ErrorHandlingUtils";
import { createUri } from "Common/UrlUtility";
import Explorer from "Explorer/Explorer";
import { useNotebook } from "Explorer/Notebook/useNotebook";
import { FeedbackParams, GenerateSQLQueryResponse } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
import { userContext } from "UserContext";
import { useQueryCopilot } from "hooks/useQueryCopilot";
@@ -30,13 +30,15 @@ export const SendQueryRequest = async ({
.getState()
.setChatMessages([...useQueryCopilot.getState().chatMessages, { source: 0, message: userPrompt }]);
try {
if (useQueryCopilot.getState().shouldAllocateContainer && userContext.features.enableCopilotPhoenixGateaway) {
await explorer.allocateContainer(PoolIdType.DefaultPoolId);
useQueryCopilot.getState().setShouldAllocateContainer(false);
if (
useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
userContext.features.enableCopilotPhoenixGateaway
) {
await explorer.allocateContainer(PoolIdType.QueryCopilot);
}
useQueryCopilot.getState().refreshCorrelationId();
const serverInfo = useNotebook.getState().notebookServerInfo;
const serverInfo = useQueryCopilot.getState().notebookServerInfo;
const queryUri = userContext.features.enableCopilotPhoenixGateaway
? createUri(serverInfo.notebookServerEndpoint, "generateSQLQuery")
@@ -58,9 +60,6 @@ export const SendQueryRequest = async ({
});
const generateSQLQueryResponse: GenerateSQLQueryResponse = await response?.json();
if (response.status === 404) {
useQueryCopilot.getState().setShouldAllocateContainer(true);
}
if (response.ok) {
if (generateSQLQueryResponse?.sql) {
let query = `Here is a query which will help you with provided prompt.\r\n **Prompt:** ${userPrompt}`;
@@ -102,7 +101,6 @@ export const SubmitFeedback = async ({
}): Promise<void> => {
try {
const { likeQuery, generatedQuery, userPrompt, description, contact } = params;
const { correlationId, shouldAllocateContainer, setShouldAllocateContainer } = useQueryCopilot();
const payload = {
containerSchema: userContext.features.enableCopilotFullSchema
? QueryCopilotSampleContainerSchema
@@ -113,25 +111,24 @@ export const SubmitFeedback = async ({
description: description || "",
contact: contact || "",
};
if (shouldAllocateContainer && userContext.features.enableCopilotPhoenixGateaway) {
await explorer.allocateContainer(PoolIdType.DefaultPoolId);
setShouldAllocateContainer(false);
if (
useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
userContext.features.enableCopilotPhoenixGateaway
) {
await explorer.allocateContainer(PoolIdType.QueryCopilot);
}
const serverInfo = useNotebook.getState().notebookServerInfo;
const serverInfo = useQueryCopilot.getState().notebookServerInfo;
const feedbackUri = userContext.features.enableCopilotPhoenixGateaway
? createUri(serverInfo.notebookServerEndpoint, "feedback")
: createUri("https://copilotorchestrater.azurewebsites.net/", "feedback");
const response = await fetch(feedbackUri, {
await fetch(feedbackUri, {
method: "POST",
headers: {
"content-type": "application/json",
"x-ms-correlationid": correlationId,
"x-ms-correlationid": useQueryCopilot.getState().correlationId,
},
body: JSON.stringify(payload),
});
if (response.status === 404) {
setShouldAllocateContainer(true);
}
} catch (error) {
handleError(error, "copilotSubmitFeedback");
}