From ae3b6eccc31a73e69dfa1e00f602475da2d68f98 Mon Sep 17 00:00:00 2001
From: Robert <robertbrignull@github.com>
Date: Fri, 7 Jul 2023 14:26:46 +0100
Subject: [PATCH] Only show AI button when available

---
 .../data-extensions-editor/DataExtensionsEditor.tsx |  1 +
 .../src/view/data-extensions-editor/LibraryRow.tsx  | 13 +++++++++----
 .../data-extensions-editor/ModeledMethodsList.tsx   |  4 ++++
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx b/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
index 5131c82e2..95f692b8d 100644
--- a/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
@@ -282,6 +282,7 @@ export function DataExtensionsEditor({
               externalApiUsages={externalApiUsages}
               unsavedModels={unsavedModels}
               modeledMethods={modeledMethods}
+              viewState={viewState}
               mode={viewState?.mode ?? Mode.Application}
               onChange={onChange}
               onSaveModelClick={onSaveModelClick}
diff --git a/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx b/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx
index 95ee11d40..620b8468f 100644
--- a/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx
@@ -13,6 +13,7 @@ import {
   VSCodeDivider,
   VSCodeTag,
 } from "@vscode/webview-ui-toolkit/react";
+import { DataExtensionEditorViewState } from "../../data-extensions-editor/shared/view-state";
 
 const LibraryContainer = styled.div`
   background-color: var(--vscode-peekViewResult-background);
@@ -69,6 +70,7 @@ type Props = {
   title: string;
   externalApiUsages: ExternalApiUsage[];
   modeledMethods: Record<string, ModeledMethod>;
+  viewState: DataExtensionEditorViewState | undefined;
   mode: Mode;
   hasUnsavedChanges: boolean;
   onChange: (
@@ -87,6 +89,7 @@ export const LibraryRow = ({
   title,
   externalApiUsages,
   modeledMethods,
+  viewState,
   mode,
   hasUnsavedChanges,
   onChange,
@@ -143,10 +146,12 @@ export const LibraryRow = ({
           </ModeledPercentage>
           {hasUnsavedChanges ? <VSCodeTag>UNSAVED</VSCodeTag> : null}
         </NameContainer>
-        <VSCodeButton appearance="icon" onClick={handleModelWithAI}>
-          <Codicon name="lightbulb-autofix" label="Model with AI" />
-          &nbsp;Model with AI
-        </VSCodeButton>
+        {viewState?.showLlmButton && (
+          <VSCodeButton appearance="icon" onClick={handleModelWithAI}>
+            <Codicon name="lightbulb-autofix" label="Model with AI" />
+            &nbsp;Model with AI
+          </VSCodeButton>
+        )}
         <VSCodeButton appearance="icon" onClick={handleModelFromSource}>
           <Codicon name="code" label="Model from source" />
           &nbsp;Model from source
diff --git a/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx b/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx
index 979c38ee7..b644af397 100644
--- a/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx
@@ -8,11 +8,13 @@ import {
   groupMethods,
   sortGroupNames,
 } from "../../data-extensions-editor/shared/sorting";
+import { DataExtensionEditorViewState } from "../../data-extensions-editor/shared/view-state";
 
 type Props = {
   externalApiUsages: ExternalApiUsage[];
   unsavedModels: Set<string>;
   modeledMethods: Record<string, ModeledMethod>;
+  viewState: DataExtensionEditorViewState | undefined;
   mode: Mode;
   onChange: (
     modelName: string,
@@ -30,6 +32,7 @@ export const ModeledMethodsList = ({
   externalApiUsages,
   unsavedModels,
   modeledMethods,
+  viewState,
   mode,
   onChange,
   onSaveModelClick,
@@ -50,6 +53,7 @@ export const ModeledMethodsList = ({
           externalApiUsages={grouped[libraryName]}
           hasUnsavedChanges={unsavedModels.has(libraryName)}
           modeledMethods={modeledMethods}
+          viewState={viewState}
           mode={mode}
           onChange={onChange}
           onSaveModelClick={onSaveModelClick}