Merge pull request #2591 from github/robertbrignull/data-llm-section

Use AI to model just a single section

Author: robertbrignull

extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx

@@ -185,14 +185,28 @@ export function DataExtensionsEditor({
     });
   }, []);
 
-  const onGenerateFromLlmClick = useCallback(() => {
+  const onGenerateAllFromLlmClick = useCallback(() => {
     vscode.postMessage({
       t: "generateExternalApiFromLlm",
       externalApiUsages,
       modeledMethods,
     });
   }, [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>
                 </>
@@ -286,6 +300,7 @@ export function DataExtensionsEditor({
               mode={viewState?.mode ?? Mode.Application}
               onChange={onChange}
               onSaveModelClick={onSaveModelClick}
+              onGenerateFromLlmClick={onGenerateFromLlmClick}
             />
           </EditorContainer>
         </>

extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx

@@ -83,6 +83,10 @@ type Props = {
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
+  onGenerateFromLlmClick: (
+    externalApiUsages: ExternalApiUsage[],
+    modeledMethods: Record<string, ModeledMethod>,
+  ) => void;
 };
 
 export const LibraryRow = ({
@@ -94,6 +98,7 @@ export const LibraryRow = ({
   hasUnsavedChanges,
   onChange,
   onSaveModelClick,
+  onGenerateFromLlmClick,
 }: Props) => {
   const modeledPercentage = useMemo(() => {
     return calculateModeledPercentage(externalApiUsages);
@@ -105,10 +110,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();

extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx

@@ -26,6 +26,10 @@ type Props = {
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
+  onGenerateFromLlmClick: (
+    externalApiUsages: ExternalApiUsage[],
+    modeledMethods: Record<string, ModeledMethod>,
+  ) => void;
 };
 
 export const ModeledMethodsList = ({
@@ -36,6 +40,7 @@ export const ModeledMethodsList = ({
   mode,
   onChange,
   onSaveModelClick,
+  onGenerateFromLlmClick,
 }: Props) => {
   const grouped = useMemo(
     () => groupMethods(externalApiUsages, mode),
@@ -57,6 +62,7 @@ export const ModeledMethodsList = ({
           mode={mode}
           onChange={onChange}
           onSaveModelClick={onSaveModelClick}
+          onGenerateFromLlmClick={onGenerateFromLlmClick}
         />
       ))}
     </>