[Query Copilot] Phoenix container implementation (#1594)

* Phoenix implementation * removing comments --------- Co-authored-by: Predrag Klepic <v-prklepic@microsoft.com>
2025-12-20 09:20:16 +00:00 · 2023-08-30 16:50:55 +02:00
parent b992742e20
commit c8e7e69aa5
6 changed files with 158 additions and 95 deletions
--- a/src/Explorer/Explorer.tsx
+++ b/src/Explorer/Explorer.tsx
@@ -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,21 +469,26 @@ 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,
-    });
-    this.notebookManager?.notebookClient
-      .getMemoryUsage()
-      .then((memoryUsageInfo) => useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo));
+    };
+    shouldUseNotebookStates
+      ? useNotebook.getState().setNotebookServerInfo(noteBookServerInfo)
+      : useQueryCopilot.getState().setNotebookServerInfo(noteBookServerInfo);
+    shouldUseNotebookStates &&
+      this.notebookManager?.notebookClient
+        .getMemoryUsage()
+        .then((memoryUsageInfo) => useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo));
  }

  public resetNotebookWorkspace(): void {
@@ -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");
--- a/src/Explorer/QueryCopilot/QueryCopilotTab.tsx
+++ b/src/Explorer/QueryCopilot/QueryCopilotTab.tsx
@@ -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`;
--- a/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.test.ts
+++ b/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.test.ts
@@ -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 () => {
--- a/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts
+++ b/src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts
@@ -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");
  }