From 48ffca31039d9d96ea932e3af8e96139abe397c1 Mon Sep 17 00:00:00 2001
From: Robert <robertbrignull@github.com>
Date: Fri, 7 Jul 2023 12:37:10 +0100
Subject: [PATCH] Use AI to model just a single section

---
 .../DataExtensionsEditor.tsx                  | 19 +++++++++++++++++--
 .../data-extensions-editor/LibraryRow.tsx     | 17 +++++++++++++----
 .../ModeledMethodsList.tsx                    |  6 ++++++
 3 files changed, 36 insertions(+), 6 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 5131c82e23c..8a5e989dc2c 100644
--- a/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
@@ -185,7 +185,7 @@ export function DataExtensionsEditor({
     });
   }, []);
 
-  const onGenerateFromLlmClick = useCallback(() => {
+  const onGenerateAllFromLlmClick = useCallback(() => {
     vscode.postMessage({
       t: "generateExternalApiFromLlm",
       externalApiUsages,
@@ -193,6 +193,20 @@ export function DataExtensionsEditor({
     });
   }, [externalApiUsages, modeledMethods]);
 
+  const onGenerateFromLlmClick = useCallback(
+    (
+      externalApiUsages: ExternalApiUsage[],
+      modeledMethods: Record<string, ModeledMethod>,
+    ) => {
+      vscode.postMessage({
+        t: "generateExternalApiFromLlm",
+        externalApiUsages,
+        modeledMethods,
+      });
+    },
+    [],
+  );
+
   const onOpenExtensionPackClick = useCallback(() => {
     vscode.postMessage({
       t: "openExtensionPack",
@@ -272,7 +286,7 @@ export function DataExtensionsEditor({
               </VSCodeButton>
               {viewState?.showLlmButton && (
                 <>
-                  <VSCodeButton onClick={onGenerateFromLlmClick}>
+                  <VSCodeButton onClick={onGenerateAllFromLlmClick}>
                     Generate using LLM
                   </VSCodeButton>
                 </>
@@ -285,6 +299,7 @@ export function DataExtensionsEditor({
               mode={viewState?.mode ?? Mode.Application}
               onChange={onChange}
               onSaveModelClick={onSaveModelClick}
+              onGenerateFromLlmClick={onGenerateFromLlmClick}
             />
           </EditorContainer>
         </>
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 95ee11d40c6..fcd98be4664 100644
--- a/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx
@@ -81,6 +81,10 @@ type Props = {
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
+  onGenerateFromLlmClick: (
+    externalApiUsages: ExternalApiUsage[],
+    modeledMethods: Record<string, ModeledMethod>,
+  ) => void;
 };
 
 export const LibraryRow = ({
@@ -91,6 +95,7 @@ export const LibraryRow = ({
   hasUnsavedChanges,
   onChange,
   onSaveModelClick,
+  onGenerateFromLlmClick,
 }: Props) => {
   const modeledPercentage = useMemo(() => {
     return calculateModeledPercentage(externalApiUsages);
@@ -102,10 +107,14 @@ export const LibraryRow = ({
     setExpanded((oldIsExpanded) => !oldIsExpanded);
   }, []);
 
-  const handleModelWithAI = useCallback(async (e: React.MouseEvent) => {
-    e.stopPropagation();
-    e.preventDefault();
-  }, []);
+  const handleModelWithAI = useCallback(
+    async (e: React.MouseEvent) => {
+      onGenerateFromLlmClick(externalApiUsages, modeledMethods);
+      e.stopPropagation();
+      e.preventDefault();
+    },
+    [externalApiUsages, modeledMethods, onGenerateFromLlmClick],
+  );
 
   const handleModelFromSource = useCallback(async (e: React.MouseEvent) => {
     e.stopPropagation();
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 979c38ee7e0..42ea7f48f2f 100644
--- a/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx
@@ -24,6 +24,10 @@ type Props = {
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
+  onGenerateFromLlmClick: (
+    externalApiUsages: ExternalApiUsage[],
+    modeledMethods: Record<string, ModeledMethod>,
+  ) => void;
 };
 
 export const ModeledMethodsList = ({
@@ -33,6 +37,7 @@ export const ModeledMethodsList = ({
   mode,
   onChange,
   onSaveModelClick,
+  onGenerateFromLlmClick,
 }: Props) => {
   const grouped = useMemo(
     () => groupMethods(externalApiUsages, mode),
@@ -53,6 +58,7 @@ export const ModeledMethodsList = ({
           mode={mode}
           onChange={onChange}
           onSaveModelClick={onSaveModelClick}
+          onGenerateFromLlmClick={onGenerateFromLlmClick}
         />
       ))}
     </>