From 5c21e9d5ee19fb3111cd11f4c585cfd172885c40 Mon Sep 17 00:00:00 2001
From: Victor Meng <vimeng@microsoft.com>
Date: Tue, 8 Aug 2023 12:21:37 -0700
Subject: [PATCH] Switch to tools federation endpoints for query copilot

---
 src/Common/Constants.ts                       | 1 +
 src/Explorer/Explorer.tsx                     | 4 +++-
 src/Explorer/QueryCopilot/QueryCopilotTab.tsx | 4 +++-
 src/Phoenix/PhoenixClient.ts                  | 9 ++++-----
 4 files changed, 11 insertions(+), 7 deletions(-)

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 e1835da4d..00332e3d5 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,
@@ -1335,5 +1335,7 @@ export default class Explorer {
 
     const sampleDataResourceTokenCollection = new ResourceTokenCollection(this, databaseId, collection, true);
     useDatabases.setState({ sampleDataResourceTokenCollection });
+
+    await this.allocateContainer();
   }
 }
diff --git a/src/Explorer/QueryCopilot/QueryCopilotTab.tsx b/src/Explorer/QueryCopilot/QueryCopilotTab.tsx
index e94b3a979..a4af8dee4 100644
--- a/src/Explorer/QueryCopilot/QueryCopilotTab.tsx
+++ b/src/Explorer/QueryCopilot/QueryCopilotTab.tsx
@@ -27,6 +27,7 @@ import { CommandButtonComponentProps } from "Explorer/Controls/CommandButton/Com
 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";
@@ -191,7 +192,8 @@ 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 response = await fetch(`${serverInfo.notebookServerEndpoint}/generateSQLQuery`, {
         method: "POST",
         headers: {
           "content-type": "application/json",
diff --git a/src/Phoenix/PhoenixClient.ts b/src/Phoenix/PhoenixClient.ts
index ddc4d563a..39b6977d5 100644
--- a/src/Phoenix/PhoenixClient.ts
+++ b/src/Phoenix/PhoenixClient.ts
@@ -1,18 +1,18 @@
 import { useDialog } from "Explorer/Controls/Dialog";
-import promiseRetry, { AbortError } from "p-retry";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { allowedJunoOrigins, validateEndpoint } from "Utils/EndpointValidation";
+import promiseRetry, { AbortError } from "p-retry";
 import {
   Areas,
   ConnectionStatusType,
   ContainerStatusType,
   HttpHeaders,
   HttpStatusCodes,
+  JunoEndpoints,
   Notebook,
 } from "../Common/Constants";
 import { getErrorMessage, getErrorStack } from "../Common/ErrorHandlingUtils";
 import * as Logger from "../Common/Logger";
-import { configContext } from "../ConfigContext";
 import {
   ContainerConnectionInfo,
   ContainerInfo,
@@ -28,7 +28,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 {
@@ -231,8 +230,8 @@ export class PhoenixClient {
       throw new Error("The Phoenix client was not initialized properly: missing ARM resourcce id");
     }
 
-    const toolsEndpoint =
-      userContext.features.phoenixEndpoint ?? userContext.features.junoEndpoint ?? configContext.JUNO_ENDPOINT;
+    const toolsEndpoint = JunoEndpoints.Test2;
+    // userContext.features.phoenixEndpoint ?? userContext.features.junoEndpoint ?? configContext.JUNO_ENDPOINT;
 
     if (!validateEndpoint(toolsEndpoint, allowedJunoOrigins)) {
       const error = `${toolsEndpoint} not allowed as tools endpoint`;