Convert yaml.ts to handle multiple models per method

This changes YAML parsing/creating functions for the model editor to
handle multiple models per method. The changes in the actual YAML
handling are fairly small because the format itself already supports
multiple models per method.

I've introduced a few helper functions to convert between the old and
new types. This should only be necessary while we're in the middle of
the transition to the new types and can be removed later. For now,
we'll just take the first model in the array when converting from the
new to the old type. This is a change in the behavior since currently
we always take the last model in the array but this behavior is
undocumented and unsupported, so it should be fine to change it.

Author: koesie10

extensions/ql-vscode/src/model-editor/auto-modeler.ts

@@ -16,6 +16,7 @@ import { QueryRunner } from "../query-server";
 import { DatabaseItem } from "../databases/local-databases";
 import { Mode } from "./shared/mode";
 import { CancellationTokenSource } from "vscode";
+import { convertToLegacyModeledMethods } from "./modeled-methods-legacy";
 
 // Limit the number of candidates we send to the model in each request
 // to avoid long requests.
@@ -192,11 +193,13 @@ export class AutoModeler {
       filename: "auto-model.yml",
     });
 
-    const loadedMethods = loadDataExtensionYaml(models);
-    if (!loadedMethods) {
+    const rawLoadedMethods = loadDataExtensionYaml(models);
+    if (!rawLoadedMethods) {
       return;
     }
 
+    const loadedMethods = convertToLegacyModeledMethods(rawLoadedMethods);
+
     // Any candidate that was part of the response is a negative result
     // meaning that the canidate is not a sink for the kinds that the LLM is checking for.
     // For now we model this as a sink neutral method, however this is subject

extensions/ql-vscode/src/model-editor/modeled-method-fs.ts

@@ -10,6 +10,11 @@ import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { load as loadYaml } from "js-yaml";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { pathsEqual } from "../common/files";
+import {
+  convertFromLegacyModeledMethods,
+  convertFromLegacyModeledMethodsFiles,
+  convertToLegacyModeledMethods,
+} from "./modeled-methods-legacy";
 
 export async function saveModeledMethods(
   extensionPack: ExtensionPack,
@@ -29,8 +34,8 @@ export async function saveModeledMethods(
   const yamls = createDataExtensionYamls(
     language,
     methods,
-    modeledMethods,
-    existingModeledMethods,
+    convertFromLegacyModeledMethods(modeledMethods),
+    convertFromLegacyModeledMethodsFiles(existingModeledMethods),
     mode,
   );
 
@@ -68,7 +73,8 @@ async function loadModeledMethodFiles(
       );
       continue;
     }
-    modeledMethodsByFile[modelFile] = modeledMethods;
+    modeledMethodsByFile[modelFile] =
+      convertToLegacyModeledMethods(modeledMethods);
   }
 
   return modeledMethodsByFile;

extensions/ql-vscode/src/model-editor/modeled-methods-legacy.ts

@@ -0,0 +1,33 @@
+import { ModeledMethod } from "./modeled-method";
+
+export function convertFromLegacyModeledMethods(
+  modeledMethods: Record<string, ModeledMethod>,
+): Record<string, ModeledMethod[]> {
+  // Convert a single ModeledMethod to an array of ModeledMethods
+  return Object.fromEntries(
+    Object.entries(modeledMethods).map(([signature, modeledMethod]) => {
+      return [signature, [modeledMethod]];
+    }),
+  );
+}
+
+export function convertToLegacyModeledMethods(
+  modeledMethods: Record<string, ModeledMethod[]>,
+): Record<string, ModeledMethod> {
+  // Always take the first modeled method in the array
+  return Object.fromEntries(
+    Object.entries(modeledMethods).map(([signature, modeledMethods]) => {
+      return [signature, modeledMethods[0]];
+    }),
+  );
+}
+
+export function convertFromLegacyModeledMethodsFiles(
+  modeledMethods: Record<string, Record<string, ModeledMethod>>,
+): Record<string, Record<string, ModeledMethod[]>> {
+  return Object.fromEntries(
+    Object.entries(modeledMethods).map(([filename, modeledMethods]) => {
+      return [filename, convertFromLegacyModeledMethods(modeledMethods)];
+    }),
+  );
+}

extensions/ql-vscode/src/model-editor/yaml.ts

@@ -71,8 +71,8 @@ ${extensions.join("\n")}`;
 export function createDataExtensionYamls(
   language: string,
   methods: Method[],
-  newModeledMethods: Record<string, ModeledMethod>,
-  existingModeledMethods: Record<string, Record<string, ModeledMethod>>,
+  newModeledMethods: Record<string, ModeledMethod[]>,
+  existingModeledMethods: Record<string, Record<string, ModeledMethod[]>>,
   mode: Mode,
 ) {
   switch (mode) {
@@ -98,11 +98,11 @@ export function createDataExtensionYamls(
 function createDataExtensionYamlsByGrouping(
   language: string,
   methods: Method[],
-  newModeledMethods: Record<string, ModeledMethod>,
-  existingModeledMethods: Record<string, Record<string, ModeledMethod>>,
+  newModeledMethods: Record<string, ModeledMethod[]>,
+  existingModeledMethods: Record<string, Record<string, ModeledMethod[]>>,
   createFilename: (method: Method) => string,
 ): Record<string, string> {
-  const methodsByFilename: Record<string, Record<string, ModeledMethod>> = {};
+  const methodsByFilename: Record<string, Record<string, ModeledMethod[]>> = {};
 
   // We only want to generate a yaml file when it's a known external API usage
   // and there are new modeled methods for it. This avoids us overwriting other
@@ -114,21 +114,25 @@ function createDataExtensionYamlsByGrouping(
   }
 
   // First populate methodsByFilename with any existing modeled methods.
-  for (const [filename, methods] of Object.entries(existingModeledMethods)) {
+  for (const [filename, methodsBySignature] of Object.entries(
+    existingModeledMethods,
+  )) {
     if (filename in methodsByFilename) {
-      for (const [signature, method] of Object.entries(methods)) {
-        methodsByFilename[filename][signature] = method;
+      for (const [signature, methods] of Object.entries(methodsBySignature)) {
+        methodsByFilename[filename][signature] = methods;
       }
     }
   }
 
   // Add the new modeled methods, potentially overwriting existing modeled methods
   // but not removing existing modeled methods that are not in the new set.
   for (const method of methods) {
-    const newMethod = newModeledMethods[method.signature];
-    if (newMethod) {
+    const newMethods = newModeledMethods[method.signature];
+    if (newMethods) {
       const filename = createFilename(method);
-      methodsByFilename[filename][newMethod.signature] = newMethod;
+
+      // Override any existing modeled methods with the new ones.
+      methodsByFilename[filename][method.signature] = newMethods;
     }
   }
 
@@ -137,7 +141,7 @@ function createDataExtensionYamlsByGrouping(
   for (const [filename, methods] of Object.entries(methodsByFilename)) {
     result[filename] = createDataExtensionYaml(
       language,
-      Object.values(methods),
+      Object.values(methods).flatMap((methods) => methods),
     );
   }
 
@@ -147,8 +151,8 @@ function createDataExtensionYamlsByGrouping(
 export function createDataExtensionYamlsForApplicationMode(
   language: string,
   methods: Method[],
-  newModeledMethods: Record<string, ModeledMethod>,
-  existingModeledMethods: Record<string, Record<string, ModeledMethod>>,
+  newModeledMethods: Record<string, ModeledMethod[]>,
+  existingModeledMethods: Record<string, Record<string, ModeledMethod[]>>,
 ): Record<string, string> {
   return createDataExtensionYamlsByGrouping(
     language,
@@ -162,8 +166,8 @@ export function createDataExtensionYamlsForApplicationMode(
 export function createDataExtensionYamlsForFrameworkMode(
   language: string,
   methods: Method[],
-  newModeledMethods: Record<string, ModeledMethod>,
-  existingModeledMethods: Record<string, Record<string, ModeledMethod>>,
+  newModeledMethods: Record<string, ModeledMethod[]>,
+  existingModeledMethods: Record<string, Record<string, ModeledMethod[]>>,
 ): Record<string, string> {
   return createDataExtensionYamlsByGrouping(
     language,
@@ -228,14 +232,14 @@ function validateModelExtensionFile(data: unknown): data is ModelExtensionFile {
 
 export function loadDataExtensionYaml(
   data: unknown,
-): Record<string, ModeledMethod> | undefined {
+): Record<string, ModeledMethod[]> | undefined {
   if (!validateModelExtensionFile(data)) {
     return undefined;
   }
 
   const extensions = data.extensions;
 
-  const modeledMethods: Record<string, ModeledMethod> = {};
+  const modeledMethods: Record<string, ModeledMethod[]> = {};
 
   for (const extension of extensions) {
     const addsTo = extension.addsTo;
@@ -250,11 +254,16 @@ export function loadDataExtensionYaml(
     }
 
     for (const row of data) {
-      const modeledMethod = definition.readModeledMethod(row);
+      const modeledMethod: ModeledMethod = definition.readModeledMethod(row);
       if (!modeledMethod) {
         continue;
       }
-      modeledMethods[modeledMethod.signature] = modeledMethod;
+
+      if (!(modeledMethod.signature in modeledMethods)) {
+        modeledMethods[modeledMethod.signature] = [];
+      }
+
+      modeledMethods[modeledMethod.signature].push(modeledMethod);
     }
   }