From 75882710adad1d623a30b50c785b7753a68b7281 Mon Sep 17 00:00:00 2001 From: Charis Kyriakou Date: Tue, 8 Aug 2023 15:52:03 +0000 Subject: [PATCH 1/2] Move automodeling logic to AutoModeler class --- .../data-extensions-editor/auto-modeler.ts | 162 ++++++++++++++++++ .../data-extensions-editor-view.ts | 149 ++-------------- 2 files changed, 181 insertions(+), 130 deletions(-) create mode 100644 extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts diff --git a/extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts b/extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts new file mode 100644 index 00000000000..d0adbbe68a1 --- /dev/null +++ b/extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts @@ -0,0 +1,162 @@ +import { ExternalApiUsage } from "./external-api-usage"; +import { ModeledMethod } from "./modeled-method"; +import { extLogger } from "../common/logging/vscode"; +import { load as loadYaml } from "js-yaml"; +import { withProgress } from "../common/vscode/progress"; +import { createAutoModelV2Request, getCandidates } from "./auto-model-v2"; +import { runAutoModelQueries } from "./auto-model-codeml-queries"; +import { loadDataExtensionYaml } from "./yaml"; +import { ModelRequest, ModelResponse, autoModelV2 } from "./auto-model-api-v2"; +import { RequestError } from "@octokit/request-error"; +import { showAndLogExceptionWithTelemetry } from "../common/logging"; +import { redactableError } from "../common/errors"; +import { App } from "../common/app"; +import { CodeQLCliServer } from "../codeql-cli/cli"; +import { QueryRunner } from "../query-server"; +import { DatabaseItem } from "../databases/local-databases"; +import { Mode } from "./shared/mode"; + +export class AutoModeler { + constructor( + private readonly app: App, + private readonly cliServer: CodeQLCliServer, + private readonly queryRunner: QueryRunner, + private readonly queryStorageDir: string, + private readonly databaseItem: DatabaseItem, + private readonly addModeledMethods: ( + modeledMethods: Record, + ) => Promise, + ) {} + + public async startModeling( + externalApiUsages: ExternalApiUsage[], + modeledMethods: Record, + mode: Mode, + ): Promise { + await this.model(externalApiUsages, modeledMethods, mode); + } + + private async model( + externalApiUsages: ExternalApiUsage[], + modeledMethods: Record, + mode: Mode, + ): Promise { + await withProgress(async (progress) => { + const maxStep = 3000; + + progress({ + step: 0, + maxStep, + message: "Retrieving usages", + }); + + // Fetch the candidates to send to the model + const candidateMethods = getCandidates( + mode, + externalApiUsages, + modeledMethods, + ); + + // If there are no candidates, there is nothing to model and we just return + if (candidateMethods.length === 0) { + void extLogger.log("No candidates to model. Stopping."); + return; + } + + const usages = await runAutoModelQueries({ + mode, + candidateMethods, + cliServer: this.cliServer, + queryRunner: this.queryRunner, + queryStorageDir: this.queryStorageDir, + databaseItem: this.databaseItem, + progress: (update) => progress({ ...update, maxStep }), + }); + if (!usages) { + return; + } + + progress({ + step: 1800, + maxStep, + message: "Creating request", + }); + + const request = await createAutoModelV2Request(mode, usages); + + progress({ + step: 2000, + maxStep, + message: "Sending request", + }); + + const response = await this.callAutoModelApi(request); + if (!response) { + return; + } + + progress({ + step: 2500, + maxStep, + message: "Parsing response", + }); + + const models = loadYaml(response.models, { + filename: "auto-model.yml", + }); + + const loadedMethods = loadDataExtensionYaml(models); + if (!loadedMethods) { + return; + } + + // 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 + // to discussion. + for (const candidate of candidateMethods) { + if (!(candidate.signature in loadedMethods)) { + loadedMethods[candidate.signature] = { + type: "neutral", + kind: "sink", + input: "", + output: "", + provenance: "ai-generated", + signature: candidate.signature, + packageName: candidate.packageName, + typeName: candidate.typeName, + methodName: candidate.methodName, + methodParameters: candidate.methodParameters, + }; + } + } + + progress({ + step: 2800, + maxStep, + message: "Applying results", + }); + + await this.addModeledMethods(loadedMethods); + }); + } + + private async callAutoModelApi( + request: ModelRequest, + ): Promise { + try { + return await autoModelV2(this.app.credentials, request); + } catch (e) { + if (e instanceof RequestError && e.status === 429) { + void showAndLogExceptionWithTelemetry( + this.app.logger, + this.app.telemetry, + redactableError(e)`Rate limit hit, please try again soon.`, + ); + return null; + } else { + throw e; + } + } + } +} diff --git a/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts b/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts index a99a0145d09..71c1a683e94 100644 --- a/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts +++ b/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts @@ -35,11 +35,6 @@ import { ExternalApiUsage } from "./external-api-usage"; import { ModeledMethod } from "./modeled-method"; import { ExtensionPack } from "./shared/extension-pack"; import { autoModel, ModelRequest, ModelResponse } from "./auto-model-api"; -import { - autoModelV2, - ModelRequest as ModelRequestV2, - ModelResponse as ModelResponseV2, -} from "./auto-model-api-v2"; import { createAutoModelRequest, parsePredictedClassifications, @@ -56,16 +51,14 @@ import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs"; import { join } from "path"; import { pickExtensionPack } from "./extension-pack-picker"; import { getLanguageDisplayName } from "../common/query-language"; -import { runAutoModelQueries } from "./auto-model-codeml-queries"; -import { createAutoModelV2Request, getCandidates } from "./auto-model-v2"; -import { load as loadYaml } from "js-yaml"; -import { loadDataExtensionYaml } from "./yaml"; -import { extLogger } from "../common/logging/vscode"; +import { AutoModeler } from "./auto-modeler"; export class DataExtensionsEditorView extends AbstractWebview< ToDataExtensionsEditorMessage, FromDataExtensionsEditorMessage > { + private readonly autoModeler: AutoModeler; + public constructor( ctx: ExtensionContext, private readonly app: App, @@ -82,6 +75,17 @@ export class DataExtensionsEditorView extends AbstractWebview< ) => void, ) { super(ctx); + + this.autoModeler = new AutoModeler( + app, + cliServer, + queryRunner, + queryStorageDir, + databaseItem, + async (modeledMethods) => { + await this.postMessage({ t: "addModeledMethods", modeledMethods }); + }, + ); } public async openView() { @@ -458,107 +462,11 @@ export class DataExtensionsEditorView extends AbstractWebview< externalApiUsages: ExternalApiUsage[], modeledMethods: Record, ): Promise { - await withProgress(async (progress) => { - const maxStep = 3000; - - progress({ - step: 0, - maxStep, - message: "Retrieving usages", - }); - - // Fetch the candidates to send to the model - const candidateMethods = getCandidates( - this.mode, - externalApiUsages, - modeledMethods, - ); - - // If there are no candidates, there is nothing to model and we just return - if (candidateMethods.length === 0) { - void extLogger.log("No candidates to model. Stopping."); - return; - } - - const usages = await runAutoModelQueries({ - mode: this.mode, - candidateMethods, - cliServer: this.cliServer, - queryRunner: this.queryRunner, - queryStorageDir: this.queryStorageDir, - databaseItem: this.databaseItem, - progress: (update) => progress({ ...update, maxStep }), - }); - if (!usages) { - return; - } - - progress({ - step: 1800, - maxStep, - message: "Creating request", - }); - - const request = await createAutoModelV2Request(this.mode, usages); - - progress({ - step: 2000, - maxStep, - message: "Sending request", - }); - - const response = await this.callAutoModelApiV2(request); - if (!response) { - return; - } - - progress({ - step: 2500, - maxStep, - message: "Parsing response", - }); - - const models = loadYaml(response.models, { - filename: "auto-model.yml", - }); - - const loadedMethods = loadDataExtensionYaml(models); - if (!loadedMethods) { - return; - } - - // 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 - // to discussion. - for (const candidate of candidateMethods) { - if (!(candidate.signature in loadedMethods)) { - loadedMethods[candidate.signature] = { - type: "neutral", - kind: "sink", - input: "", - output: "", - provenance: "ai-generated", - signature: candidate.signature, - packageName: candidate.packageName, - typeName: candidate.typeName, - methodName: candidate.methodName, - methodParameters: candidate.methodParameters, - }; - } - } - - progress({ - step: 2800, - maxStep, - message: "Applying results", - }); - - await this.postMessage({ - t: "addModeledMethods", - modeledMethods: loadedMethods, - }); - }); + await this.autoModeler.startModeling( + externalApiUsages, + modeledMethods, + this.mode, + ); } private async modelDependency(): Promise { @@ -688,23 +596,4 @@ export class DataExtensionsEditorView extends AbstractWebview< } } } - - private async callAutoModelApiV2( - request: ModelRequestV2, - ): Promise { - try { - return await autoModelV2(this.app.credentials, request); - } catch (e) { - if (e instanceof RequestError && e.status === 429) { - void showAndLogExceptionWithTelemetry( - this.app.logger, - this.app.telemetry, - redactableError(e)`Rate limit hit, please try again soon.`, - ); - return null; - } else { - throw e; - } - } - } } From 0cc8e681556db40ce84c34318617586836dd8e87 Mon Sep 17 00:00:00 2001 From: Charis Kyriakou Date: Tue, 8 Aug 2023 15:56:04 +0000 Subject: [PATCH 2/2] Extract modeling of specific set of candidates into its own function --- .../data-extensions-editor/auto-modeler.ts | 147 ++++++++++-------- 1 file changed, 78 insertions(+), 69 deletions(-) diff --git a/extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts b/extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts index d0adbbe68a1..1a82b2cf32f 100644 --- a/extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts +++ b/extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts @@ -1,8 +1,8 @@ -import { ExternalApiUsage } from "./external-api-usage"; +import { ExternalApiUsage, MethodSignature } from "./external-api-usage"; import { ModeledMethod } from "./modeled-method"; import { extLogger } from "../common/logging/vscode"; import { load as loadYaml } from "js-yaml"; -import { withProgress } from "../common/vscode/progress"; +import { ProgressCallback, withProgress } from "../common/vscode/progress"; import { createAutoModelV2Request, getCandidates } from "./auto-model-v2"; import { runAutoModelQueries } from "./auto-model-codeml-queries"; import { loadDataExtensionYaml } from "./yaml"; @@ -33,10 +33,10 @@ export class AutoModeler { modeledMethods: Record, mode: Mode, ): Promise { - await this.model(externalApiUsages, modeledMethods, mode); + await this.modelDependency(externalApiUsages, modeledMethods, mode); } - private async model( + private async modelDependency( externalApiUsages: ExternalApiUsage[], modeledMethods: Record, mode: Mode, @@ -63,82 +63,91 @@ export class AutoModeler { return; } - const usages = await runAutoModelQueries({ - mode, - candidateMethods, - cliServer: this.cliServer, - queryRunner: this.queryRunner, - queryStorageDir: this.queryStorageDir, - databaseItem: this.databaseItem, - progress: (update) => progress({ ...update, maxStep }), - }); - if (!usages) { - return; - } + await this.modelCandidates(candidateMethods, mode, progress, maxStep); + }); + } - progress({ - step: 1800, - maxStep, - message: "Creating request", - }); + private async modelCandidates( + candidateMethods: MethodSignature[], + mode: Mode, + progress: ProgressCallback, + maxStep: number, + ): Promise { + const usages = await runAutoModelQueries({ + mode, + candidateMethods, + cliServer: this.cliServer, + queryRunner: this.queryRunner, + queryStorageDir: this.queryStorageDir, + databaseItem: this.databaseItem, + progress: (update) => progress({ ...update, maxStep }), + }); + if (!usages) { + return; + } - const request = await createAutoModelV2Request(mode, usages); + progress({ + step: 1800, + maxStep, + message: "Creating request", + }); - progress({ - step: 2000, - maxStep, - message: "Sending request", - }); + const request = await createAutoModelV2Request(mode, usages); - const response = await this.callAutoModelApi(request); - if (!response) { - return; - } + progress({ + step: 2000, + maxStep, + message: "Sending request", + }); - progress({ - step: 2500, - maxStep, - message: "Parsing response", - }); + const response = await this.callAutoModelApi(request); + if (!response) { + return; + } - const models = loadYaml(response.models, { - filename: "auto-model.yml", - }); + progress({ + step: 2500, + maxStep, + message: "Parsing response", + }); - const loadedMethods = loadDataExtensionYaml(models); - if (!loadedMethods) { - return; - } + const models = loadYaml(response.models, { + filename: "auto-model.yml", + }); - // 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 - // to discussion. - for (const candidate of candidateMethods) { - if (!(candidate.signature in loadedMethods)) { - loadedMethods[candidate.signature] = { - type: "neutral", - kind: "sink", - input: "", - output: "", - provenance: "ai-generated", - signature: candidate.signature, - packageName: candidate.packageName, - typeName: candidate.typeName, - methodName: candidate.methodName, - methodParameters: candidate.methodParameters, - }; - } - } + const loadedMethods = loadDataExtensionYaml(models); + if (!loadedMethods) { + return; + } - progress({ - step: 2800, - maxStep, - message: "Applying results", - }); + // 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 + // to discussion. + for (const candidate of candidateMethods) { + if (!(candidate.signature in loadedMethods)) { + loadedMethods[candidate.signature] = { + type: "neutral", + kind: "sink", + input: "", + output: "", + provenance: "ai-generated", + signature: candidate.signature, + packageName: candidate.packageName, + typeName: candidate.typeName, + methodName: candidate.methodName, + methodParameters: candidate.methodParameters, + }; + } + } - await this.addModeledMethods(loadedMethods); + progress({ + step: 2800, + maxStep, + message: "Applying results", }); + + await this.addModeledMethods(loadedMethods); } private async callAutoModelApi(