mirror of
https://github.com/Azure/cosmos-explorer.git
synced 2024-11-25 15:06:55 +00:00
[Query Copilot] Phoenix container implementation (#1594)
* Phoenix implementation * removing comments --------- Co-authored-by: Predrag Klepic <v-prklepic@microsoft.com>
This commit is contained in:
Predrag Klepic
GitHub
parent
b992742e20
commit
c8e7e69aa5
@ -5,6 +5,7 @@ import { Platform } from "ConfigContext";
|
||||
import { MessageTypes } from "Contracts/ExplorerContracts";
|
||||
import { IGalleryItem } from "Juno/JunoClient";
|
||||
import { allowedNotebookServerUrls, validateEndpoint } from "Utils/EndpointValidation";
|
||||
import { useQueryCopilot } from "hooks/useQueryCopilot";
|
||||
import * as ko from "knockout";
|
||||
import React from "react";
|
||||
import _ from "underscore";
|
||||
@ -386,8 +387,14 @@ export default class Explorer {
|
||||
}
|
||||
|
||||
public async allocateContainer(poolId: PoolIdType): Promise<void> {
|
||||
const notebookServerInfo = useNotebook.getState().notebookServerInfo;
|
||||
const isAllocating = useNotebook.getState().isAllocating;
|
||||
const shouldUseNotebookStates = poolId === PoolIdType.DefaultPoolId ? true : false;
|
||||
const notebookServerInfo = shouldUseNotebookStates
|
||||
? useNotebook.getState().notebookServerInfo
|
||||
: useQueryCopilot.getState().notebookServerInfo;
|
||||
|
||||
const isAllocating = shouldUseNotebookStates
|
||||
? useNotebook.getState().isAllocating
|
||||
: useQueryCopilot.getState().isAllocatingContainer;
|
||||
if (
|
||||
isAllocating === false &&
|
||||
(notebookServerInfo === undefined ||
|
||||
@ -395,23 +402,28 @@ export default class Explorer {
|
||||
) {
|
||||
const provisionData: IProvisionData = {
|
||||
cosmosEndpoint: userContext?.databaseAccount?.properties?.documentEndpoint,
|
||||
poolId: poolId === PoolIdType.DefaultPoolId ? undefined : poolId,
|
||||
poolId: shouldUseNotebookStates ? undefined : poolId,
|
||||
};
|
||||
const connectionStatus: ContainerConnectionInfo = {
|
||||
status: ConnectionStatusType.Connecting,
|
||||
};
|
||||
useNotebook.getState().setConnectionInfo(connectionStatus);
|
||||
|
||||
shouldUseNotebookStates && useNotebook.getState().setConnectionInfo(connectionStatus);
|
||||
|
||||
let connectionInfo;
|
||||
try {
|
||||
TelemetryProcessor.traceStart(Action.PhoenixConnection, {
|
||||
dataExplorerArea: Areas.Notebook,
|
||||
});
|
||||
useNotebook.getState().setIsAllocating(true);
|
||||
shouldUseNotebookStates
|
||||
? useNotebook.getState().setIsAllocating(true)
|
||||
: useQueryCopilot.getState().setIsAllocatingContainer(true);
|
||||
|
||||
connectionInfo = await this.phoenixClient.allocateContainer(provisionData);
|
||||
if (!connectionInfo?.data?.phoenixServiceUrl) {
|
||||
throw new Error(`PhoenixServiceUrl is invalid!`);
|
||||
}
|
||||
await this.setNotebookInfo(connectionInfo, connectionStatus);
|
||||
await this.setNotebookInfo(shouldUseNotebookStates, connectionInfo, connectionStatus);
|
||||
TelemetryProcessor.traceSuccess(Action.PhoenixConnection, {
|
||||
dataExplorerArea: Areas.Notebook,
|
||||
});
|
||||
@ -423,7 +435,9 @@ export default class Explorer {
|
||||
errorStack: getErrorStack(error),
|
||||
});
|
||||
connectionStatus.status = ConnectionStatusType.Failed;
|
||||
useNotebook.getState().resetContainerConnection(connectionStatus);
|
||||
shouldUseNotebookStates
|
||||
? useNotebook.getState().resetContainerConnection(connectionStatus)
|
||||
: useQueryCopilot.getState().resetContainerConnection();
|
||||
if (error?.status === HttpStatusCodes.Forbidden && error.message) {
|
||||
useDialog.getState().showOkModalDialog("Connection Failed", `${error.message}`);
|
||||
} else {
|
||||
@ -436,7 +450,9 @@ export default class Explorer {
|
||||
}
|
||||
throw error;
|
||||
} finally {
|
||||
useNotebook.getState().setIsAllocating(false);
|
||||
shouldUseNotebookStates
|
||||
? useNotebook.getState().setIsAllocating(false)
|
||||
: useQueryCopilot.getState().setIsAllocatingContainer(false);
|
||||
this.refreshCommandBarButtons();
|
||||
this.refreshNotebookList();
|
||||
this._isInitializingNotebooks = false;
|
||||
@ -445,6 +461,7 @@ export default class Explorer {
|
||||
}
|
||||
|
||||
private async setNotebookInfo(
|
||||
shouldUseNotebookStates: boolean,
|
||||
connectionInfo: IResponse<IPhoenixServiceInfo>,
|
||||
connectionStatus: DataModels.ContainerConnectionInfo
|
||||
) {
|
||||
@ -452,18 +469,23 @@ export default class Explorer {
|
||||
forwardingId: connectionInfo.data.forwardingId,
|
||||
dbAccountName: userContext.databaseAccount.name,
|
||||
};
|
||||
await this.phoenixClient.initiateContainerHeartBeat(containerData);
|
||||
await this.phoenixClient.initiateContainerHeartBeat(shouldUseNotebookStates, containerData);
|
||||
|
||||
connectionStatus.status = ConnectionStatusType.Connected;
|
||||
useNotebook.getState().setConnectionInfo(connectionStatus);
|
||||
useNotebook.getState().setNotebookServerInfo({
|
||||
shouldUseNotebookStates && useNotebook.getState().setConnectionInfo(connectionStatus);
|
||||
|
||||
const noteBookServerInfo = {
|
||||
notebookServerEndpoint:
|
||||
(validateEndpoint(userContext.features.notebookServerUrl, allowedNotebookServerUrls) &&
|
||||
userContext.features.notebookServerUrl) ||
|
||||
connectionInfo.data.phoenixServiceUrl,
|
||||
authToken: userContext.features.notebookServerToken || connectionInfo.data.authToken,
|
||||
forwardingId: connectionInfo.data.forwardingId,
|
||||
});
|
||||
};
|
||||
shouldUseNotebookStates
|
||||
? useNotebook.getState().setNotebookServerInfo(noteBookServerInfo)
|
||||
: useQueryCopilot.getState().setNotebookServerInfo(noteBookServerInfo);
|
||||
shouldUseNotebookStates &&
|
||||
this.notebookManager?.notebookClient
|
||||
.getMemoryUsage()
|
||||
.then((memoryUsageInfo) => useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo));
|
||||
@ -540,7 +562,7 @@ export default class Explorer {
|
||||
throw new Error(`Reset Workspace: PhoenixServiceUrl is invalid!`);
|
||||
}
|
||||
if (useNotebook.getState().isPhoenixNotebooks) {
|
||||
await this.setNotebookInfo(connectionInfo, connectionStatus);
|
||||
await this.setNotebookInfo(true, connectionInfo, connectionStatus);
|
||||
useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
|
||||
}
|
||||
logConsoleInfo("Successfully reset notebook workspace");
|
||||
|
||||
@ -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`;
|
||||
|
||||
@ -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(() => {
|
||||
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 () => {
|
||||
|
||||
@ -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");
|
||||
}
|
||||
|
||||
@ -3,6 +3,7 @@ import { useDialog } from "Explorer/Controls/Dialog";
|
||||
import { Action } from "Shared/Telemetry/TelemetryConstants";
|
||||
import { userContext } from "UserContext";
|
||||
import { allowedJunoOrigins, validateEndpoint } from "Utils/EndpointValidation";
|
||||
import { useQueryCopilot } from "hooks/useQueryCopilot";
|
||||
import promiseRetry, { AbortError } from "p-retry";
|
||||
import {
|
||||
Areas,
|
||||
@ -96,20 +97,27 @@ export class PhoenixClient {
|
||||
}
|
||||
}
|
||||
|
||||
public async initiateContainerHeartBeat(containerData: IContainerData) {
|
||||
public async initiateContainerHeartBeat(shouldUseNotebookStates: boolean, containerData: IContainerData) {
|
||||
if (this.containerHealthHandler) {
|
||||
clearTimeout(this.containerHealthHandler);
|
||||
}
|
||||
await this.getContainerHealth(Notebook.containerStatusHeartbeatDelayMs, containerData);
|
||||
await this.getContainerHealth(shouldUseNotebookStates, Notebook.containerStatusHeartbeatDelayMs, containerData);
|
||||
}
|
||||
|
||||
private scheduleContainerHeartbeat(delayMs: number, containerData: IContainerData): void {
|
||||
private scheduleContainerHeartbeat(
|
||||
shouldUseNotebookStates: boolean,
|
||||
delayMs: number,
|
||||
containerData: IContainerData
|
||||
): void {
|
||||
this.containerHealthHandler = setTimeout(async () => {
|
||||
await this.getContainerHealth(delayMs, containerData);
|
||||
await this.getContainerHealth(shouldUseNotebookStates, delayMs, containerData);
|
||||
}, delayMs);
|
||||
}
|
||||
|
||||
private async getContainerStatusAsync(containerData: IContainerData): Promise<ContainerInfo> {
|
||||
private async getContainerStatusAsync(
|
||||
shouldUseNotebookStates: boolean,
|
||||
containerData: IContainerData
|
||||
): Promise<ContainerInfo> {
|
||||
try {
|
||||
const runContainerStatusAsync = async () => {
|
||||
const response = await window.fetch(
|
||||
@ -136,8 +144,11 @@ export class PhoenixClient {
|
||||
dataExplorerArea: Areas.Notebook,
|
||||
message: getErrorMessage(error),
|
||||
});
|
||||
useNotebook.getState().resetContainerConnection(connectionStatus);
|
||||
useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
|
||||
shouldUseNotebookStates
|
||||
? useNotebook.getState().resetContainerConnection(connectionStatus)
|
||||
: useQueryCopilot.getState().resetContainerConnection();
|
||||
shouldUseNotebookStates && useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
|
||||
shouldUseNotebookStates &&
|
||||
useDialog
|
||||
.getState()
|
||||
.showOkModalDialog(
|
||||
@ -163,8 +174,10 @@ export class PhoenixClient {
|
||||
const connectionStatus: ContainerConnectionInfo = {
|
||||
status: ConnectionStatusType.Failed,
|
||||
};
|
||||
useNotebook.getState().resetContainerConnection(connectionStatus);
|
||||
useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
|
||||
shouldUseNotebookStates
|
||||
? useNotebook.getState().resetContainerConnection(connectionStatus)
|
||||
: useQueryCopilot.getState().resetContainerConnection();
|
||||
shouldUseNotebookStates && useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
|
||||
return {
|
||||
durationLeftInMinutes: undefined,
|
||||
phoenixServerInfo: undefined,
|
||||
@ -173,11 +186,17 @@ export class PhoenixClient {
|
||||
}
|
||||
}
|
||||
|
||||
private async getContainerHealth(delayMs: number, containerData: IContainerData) {
|
||||
const containerInfo = await this.getContainerStatusAsync(containerData);
|
||||
useNotebook.getState().setContainerStatus(containerInfo);
|
||||
if (useNotebook.getState().containerStatus?.status === ContainerStatusType.Active) {
|
||||
this.scheduleContainerHeartbeat(delayMs, containerData);
|
||||
private async getContainerHealth(shouldUseNotebookStates: boolean, delayMs: number, containerData: IContainerData) {
|
||||
const containerInfo = await this.getContainerStatusAsync(shouldUseNotebookStates, containerData);
|
||||
shouldUseNotebookStates
|
||||
? useNotebook.getState().setContainerStatus(containerInfo)
|
||||
: useQueryCopilot.getState().setContainerStatus(containerInfo);
|
||||
|
||||
const containerStatus = shouldUseNotebookStates
|
||||
? useNotebook.getState().containerStatus?.status
|
||||
: useQueryCopilot.getState().containerStatus?.status;
|
||||
if (containerStatus === ContainerStatusType.Active) {
|
||||
this.scheduleContainerHeartbeat(shouldUseNotebookStates, delayMs, containerData);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -2,7 +2,10 @@ import { MinimalQueryIterator } from "Common/IteratorUtilities";
|
||||
import { QueryResults } from "Contracts/ViewModels";
|
||||
import { CopilotMessage } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
|
||||
import { guid } from "Explorer/Tables/Utilities";
|
||||
import { useTabs } from "hooks/useTabs";
|
||||
import create, { UseStore } from "zustand";
|
||||
import * as DataModels from "../Contracts/DataModels";
|
||||
import { ContainerInfo } from "../Contracts/DataModels";
|
||||
|
||||
export interface QueryCopilotState {
|
||||
generatedQuery: string;
|
||||
@ -32,10 +35,12 @@ export interface QueryCopilotState {
|
||||
showWelcomeSidebar: boolean;
|
||||
showCopilotSidebar: boolean;
|
||||
chatMessages: CopilotMessage[];
|
||||
shouldAllocateContainer: boolean;
|
||||
shouldIncludeInMessages: boolean;
|
||||
showExplanationBubble: boolean;
|
||||
showQueryExplanation: boolean;
|
||||
notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo;
|
||||
containerStatus: ContainerInfo;
|
||||
isAllocatingContainer: boolean;
|
||||
|
||||
openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void;
|
||||
closeFeedbackModal: () => void;
|
||||
@ -65,11 +70,14 @@ export interface QueryCopilotState {
|
||||
setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => void;
|
||||
setShowCopilotSidebar: (showCopilotSidebar: boolean) => void;
|
||||
setChatMessages: (chatMessages: CopilotMessage[]) => void;
|
||||
setShouldAllocateContainer: (shouldAllocateContainer: boolean) => void;
|
||||
setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => void;
|
||||
setShowExplanationBubble: (showExplanationBubble: boolean) => void;
|
||||
setShowQueryExplanation: (showQueryExplanation: boolean) => void;
|
||||
setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => void;
|
||||
setContainerStatus: (containerStatus: ContainerInfo) => void;
|
||||
setIsAllocatingContainer: (isAllocatingContainer: boolean) => void;
|
||||
|
||||
resetContainerConnection: () => void;
|
||||
resetQueryCopilotStates: () => void;
|
||||
}
|
||||
|
||||
@ -103,10 +111,20 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
|
||||
showWelcomeSidebar: true,
|
||||
showCopilotSidebar: false,
|
||||
chatMessages: [],
|
||||
shouldAllocateContainer: true,
|
||||
shouldIncludeInMessages: true,
|
||||
showExplanationBubble: false,
|
||||
showQueryExplanation: false,
|
||||
notebookServerInfo: {
|
||||
notebookServerEndpoint: undefined,
|
||||
authToken: undefined,
|
||||
forwardingId: undefined,
|
||||
},
|
||||
containerStatus: {
|
||||
status: undefined,
|
||||
durationLeftInMinutes: undefined,
|
||||
phoenixServerInfo: undefined,
|
||||
},
|
||||
isAllocatingContainer: false,
|
||||
|
||||
openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
|
||||
set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
|
||||
@ -138,10 +156,24 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
|
||||
setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => set({ showWelcomeSidebar }),
|
||||
setShowCopilotSidebar: (showCopilotSidebar: boolean) => set({ showCopilotSidebar }),
|
||||
setChatMessages: (chatMessages: CopilotMessage[]) => set({ chatMessages }),
|
||||
setShouldAllocateContainer: (shouldAllocateContainer: boolean) => set({ shouldAllocateContainer }),
|
||||
setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => set({ shouldIncludeInMessages }),
|
||||
setShowExplanationBubble: (showExplanationBubble: boolean) => set({ showExplanationBubble }),
|
||||
setShowQueryExplanation: (showQueryExplanation: boolean) => set({ showQueryExplanation }),
|
||||
setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
|
||||
set({ notebookServerInfo }),
|
||||
setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }),
|
||||
setIsAllocatingContainer: (isAllocatingContainer: boolean) => set({ isAllocatingContainer }),
|
||||
|
||||
resetContainerConnection: (): void => {
|
||||
useTabs.getState().closeAllNotebookTabs(true);
|
||||
useQueryCopilot.getState().setNotebookServerInfo(undefined);
|
||||
useQueryCopilot.getState().setIsAllocatingContainer(false);
|
||||
useQueryCopilot.getState().setContainerStatus({
|
||||
status: undefined,
|
||||
durationLeftInMinutes: undefined,
|
||||
phoenixServerInfo: undefined,
|
||||
});
|
||||
},
|
||||
|
||||
resetQueryCopilotStates: () => {
|
||||
set((state) => ({
|
||||
@ -172,10 +204,20 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
|
||||
wasCopilotUsed: false,
|
||||
showCopilotSidebar: false,
|
||||
chatMessages: [],
|
||||
shouldAllocateContainer: true,
|
||||
shouldIncludeInMessages: true,
|
||||
showExplanationBubble: false,
|
||||
showQueryExplanation: false,
|
||||
notebookServerInfo: {
|
||||
notebookServerEndpoint: undefined,
|
||||
authToken: undefined,
|
||||
forwardingId: undefined,
|
||||
},
|
||||
containerStatus: {
|
||||
status: undefined,
|
||||
durationLeftInMinutes: undefined,
|
||||
phoenixServerInfo: undefined,
|
||||
},
|
||||
isAllocatingContainer: false,
|
||||
}));
|
||||
},
|
||||
}));
|
||||
|
||||
Loading…
Reference in New Issue
Block a user