Convert yaml.ts to handle multiple models per method

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);
     }
   }