Skip to content

Commit 6dc0b14

Browse files
chariskcharisk
authored
Merge pull request #2680 from github/charisk/move-automodeling
Move automodeling logic out of view
2 parents f52a512 + 0cc8e68 commit 6dc0b14

File tree

2 files changed

+190
-130
lines changed

2 files changed

+190
-130
lines changed

‎extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
import { ExternalApiUsage, MethodSignature } from "./external-api-usage";
2+
import { ModeledMethod } from "./modeled-method";
3+
import { extLogger } from "../common/logging/vscode";
4+
import { load as loadYaml } from "js-yaml";
5+
import { ProgressCallback, withProgress } from "../common/vscode/progress";
6+
import { createAutoModelV2Request, getCandidates } from "./auto-model-v2";
7+
import { runAutoModelQueries } from "./auto-model-codeml-queries";
8+
import { loadDataExtensionYaml } from "./yaml";
9+
import { ModelRequest, ModelResponse, autoModelV2 } from "./auto-model-api-v2";
10+
import { RequestError } from "@octokit/request-error";
11+
import { showAndLogExceptionWithTelemetry } from "../common/logging";
12+
import { redactableError } from "../common/errors";
13+
import { App } from "../common/app";
14+
import { CodeQLCliServer } from "../codeql-cli/cli";
15+
import { QueryRunner } from "../query-server";
16+
import { DatabaseItem } from "../databases/local-databases";
17+
import { Mode } from "./shared/mode";
18+
19+
export class AutoModeler {
20+
constructor(
21+
private readonly app: App,
22+
private readonly cliServer: CodeQLCliServer,
23+
private readonly queryRunner: QueryRunner,
24+
private readonly queryStorageDir: string,
25+
private readonly databaseItem: DatabaseItem,
26+
private readonly addModeledMethods: (
27+
modeledMethods: Record<string, ModeledMethod>,
28+
) => Promise<void>,
29+
) {}
30+
31+
public async startModeling(
32+
externalApiUsages: ExternalApiUsage[],
33+
modeledMethods: Record<string, ModeledMethod>,
34+
mode: Mode,
35+
): Promise<void> {
36+
await this.modelDependency(externalApiUsages, modeledMethods, mode);
37+
}
38+
39+
private async modelDependency(
40+
externalApiUsages: ExternalApiUsage[],
41+
modeledMethods: Record<string, ModeledMethod>,
42+
mode: Mode,
43+
): Promise<void> {
44+
await withProgress(async (progress) => {
45+
const maxStep = 3000;
46+
47+
progress({
48+
step: 0,
49+
maxStep,
50+
message: "Retrieving usages",
51+
});
52+
53+
// Fetch the candidates to send to the model
54+
const candidateMethods = getCandidates(
55+
mode,
56+
externalApiUsages,
57+
modeledMethods,
58+
);
59+
60+
// If there are no candidates, there is nothing to model and we just return
61+
if (candidateMethods.length === 0) {
62+
void extLogger.log("No candidates to model. Stopping.");
63+
return;
64+
}
65+
66+
await this.modelCandidates(candidateMethods, mode, progress, maxStep);
67+
});
68+
}
69+
70+
private async modelCandidates(
71+
candidateMethods: MethodSignature[],
72+
mode: Mode,
73+
progress: ProgressCallback,
74+
maxStep: number,
75+
): Promise<void> {
76+
const usages = await runAutoModelQueries({
77+
mode,
78+
candidateMethods,
79+
cliServer: this.cliServer,
80+
queryRunner: this.queryRunner,
81+
queryStorageDir: this.queryStorageDir,
82+
databaseItem: this.databaseItem,
83+
progress: (update) => progress({ ...update, maxStep }),
84+
});
85+
if (!usages) {
86+
return;
87+
}
88+
89+
progress({
90+
step: 1800,
91+
maxStep,
92+
message: "Creating request",
93+
});
94+
95+
const request = await createAutoModelV2Request(mode, usages);
96+
97+
progress({
98+
step: 2000,
99+
maxStep,
100+
message: "Sending request",
101+
});
102+
103+
const response = await this.callAutoModelApi(request);
104+
if (!response) {
105+
return;
106+
}
107+
108+
progress({
109+
step: 2500,
110+
maxStep,
111+
message: "Parsing response",
112+
});
113+
114+
const models = loadYaml(response.models, {
115+
filename: "auto-model.yml",
116+
});
117+
118+
const loadedMethods = loadDataExtensionYaml(models);
119+
if (!loadedMethods) {
120+
return;
121+
}
122+
123+
// Any candidate that was part of the response is a negative result
124+
// meaning that the canidate is not a sink for the kinds that the LLM is checking for.
125+
// For now we model this as a sink neutral method, however this is subject
126+
// to discussion.
127+
for (const candidate of candidateMethods) {
128+
if (!(candidate.signature in loadedMethods)) {
129+
loadedMethods[candidate.signature] = {
130+
type: "neutral",
131+
kind: "sink",
132+
input: "",
133+
output: "",
134+
provenance: "ai-generated",
135+
signature: candidate.signature,
136+
packageName: candidate.packageName,
137+
typeName: candidate.typeName,
138+
methodName: candidate.methodName,
139+
methodParameters: candidate.methodParameters,
140+
};
141+
}
142+
}
143+
144+
progress({
145+
step: 2800,
146+
maxStep,
147+
message: "Applying results",
148+
});
149+
150+
await this.addModeledMethods(loadedMethods);
151+
}
152+
153+
private async callAutoModelApi(
154+
request: ModelRequest,
155+
): Promise<ModelResponse | null> {
156+
try {
157+
return await autoModelV2(this.app.credentials, request);
158+
} catch (e) {
159+
if (e instanceof RequestError && e.status === 429) {
160+
void showAndLogExceptionWithTelemetry(
161+
this.app.logger,
162+
this.app.telemetry,
163+
redactableError(e)`Rate limit hit, please try again soon.`,
164+
);
165+
return null;
166+
} else {
167+
throw e;
168+
}
169+
}
170+
}
171+
}

‎extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts‎

Lines changed: 19 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -35,11 +35,6 @@ import { ExternalApiUsage } from "./external-api-usage";
3535
import { ModeledMethod } from "./modeled-method";
3636
import { ExtensionPack } from "./shared/extension-pack";
3737
import { autoModel, ModelRequest, ModelResponse } from "./auto-model-api";
38-
import {
39-
autoModelV2,
40-
ModelRequest as ModelRequestV2,
41-
ModelResponse as ModelResponseV2,
42-
} from "./auto-model-api-v2";
4338
import {
4439
createAutoModelRequest,
4540
parsePredictedClassifications,
@@ -56,16 +51,14 @@ import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
5651
import { join } from "path";
5752
import { pickExtensionPack } from "./extension-pack-picker";
5853
import { getLanguageDisplayName } from "../common/query-language";
59-
import { runAutoModelQueries } from "./auto-model-codeml-queries";
60-
import { createAutoModelV2Request, getCandidates } from "./auto-model-v2";
61-
import { load as loadYaml } from "js-yaml";
62-
import { loadDataExtensionYaml } from "./yaml";
63-
import { extLogger } from "../common/logging/vscode";
54+
import { AutoModeler } from "./auto-modeler";
6455

6556
export class DataExtensionsEditorView extends AbstractWebview<
6657
ToDataExtensionsEditorMessage,
6758
FromDataExtensionsEditorMessage
6859
> {
60+
private readonly autoModeler: AutoModeler;
61+
6962
public constructor(
7063
ctx: ExtensionContext,
7164
private readonly app: App,
@@ -82,6 +75,17 @@ export class DataExtensionsEditorView extends AbstractWebview<
8275
) => void,
8376
) {
8477
super(ctx);
78+
79+
this.autoModeler = new AutoModeler(
80+
app,
81+
cliServer,
82+
queryRunner,
83+
queryStorageDir,
84+
databaseItem,
85+
async (modeledMethods) => {
86+
await this.postMessage({ t: "addModeledMethods", modeledMethods });
87+
},
88+
);
8589
}
8690

8791
public async openView() {
@@ -458,107 +462,11 @@ export class DataExtensionsEditorView extends AbstractWebview<
458462
externalApiUsages: ExternalApiUsage[],
459463
modeledMethods: Record<string, ModeledMethod>,
460464
): Promise<void> {
461-
await withProgress(async (progress) => {
462-
const maxStep = 3000;
463-
464-
progress({
465-
step: 0,
466-
maxStep,
467-
message: "Retrieving usages",
468-
});
469-
470-
// Fetch the candidates to send to the model
471-
const candidateMethods = getCandidates(
472-
this.mode,
473-
externalApiUsages,
474-
modeledMethods,
475-
);
476-
477-
// If there are no candidates, there is nothing to model and we just return
478-
if (candidateMethods.length === 0) {
479-
void extLogger.log("No candidates to model. Stopping.");
480-
return;
481-
}
482-
483-
const usages = await runAutoModelQueries({
484-
mode: this.mode,
485-
candidateMethods,
486-
cliServer: this.cliServer,
487-
queryRunner: this.queryRunner,
488-
queryStorageDir: this.queryStorageDir,
489-
databaseItem: this.databaseItem,
490-
progress: (update) => progress({ ...update, maxStep }),
491-
});
492-
if (!usages) {
493-
return;
494-
}
495-
496-
progress({
497-
step: 1800,
498-
maxStep,
499-
message: "Creating request",
500-
});
501-
502-
const request = await createAutoModelV2Request(this.mode, usages);
503-
504-
progress({
505-
step: 2000,
506-
maxStep,
507-
message: "Sending request",
508-
});
509-
510-
const response = await this.callAutoModelApiV2(request);
511-
if (!response) {
512-
return;
513-
}
514-
515-
progress({
516-
step: 2500,
517-
maxStep,
518-
message: "Parsing response",
519-
});
520-
521-
const models = loadYaml(response.models, {
522-
filename: "auto-model.yml",
523-
});
524-
525-
const loadedMethods = loadDataExtensionYaml(models);
526-
if (!loadedMethods) {
527-
return;
528-
}
529-
530-
// Any candidate that was part of the response is a negative result
531-
// meaning that the canidate is not a sink for the kinds that the LLM is checking for.
532-
// For now we model this as a sink neutral method, however this is subject
533-
// to discussion.
534-
for (const candidate of candidateMethods) {
535-
if (!(candidate.signature in loadedMethods)) {
536-
loadedMethods[candidate.signature] = {
537-
type: "neutral",
538-
kind: "sink",
539-
input: "",
540-
output: "",
541-
provenance: "ai-generated",
542-
signature: candidate.signature,
543-
packageName: candidate.packageName,
544-
typeName: candidate.typeName,
545-
methodName: candidate.methodName,
546-
methodParameters: candidate.methodParameters,
547-
};
548-
}
549-
}
550-
551-
progress({
552-
step: 2800,
553-
maxStep,
554-
message: "Applying results",
555-
});
556-
557-
await this.postMessage({
558-
t: "addModeledMethods",
559-
modeledMethods: loadedMethods,
560-
});
561-
});
465+
await this.autoModeler.startModeling(
466+
externalApiUsages,
467+
modeledMethods,
468+
this.mode,
469+
);
562470
}
563471

564472
private async modelDependency(): Promise<void> {
@@ -688,23 +596,4 @@ export class DataExtensionsEditorView extends AbstractWebview<
688596
}
689597
}
690598
}
691-
692-
private async callAutoModelApiV2(
693-
request: ModelRequestV2,
694-
): Promise<ModelResponseV2 | null> {
695-
try {
696-
return await autoModelV2(this.app.credentials, request);
697-
} catch (e) {
698-
if (e instanceof RequestError && e.status === 429) {
699-
void showAndLogExceptionWithTelemetry(
700-
this.app.logger,
701-
this.app.telemetry,
702-
redactableError(e)`Rate limit hit, please try again soon.`,
703-
);
704-
return null;
705-
} else {
706-
throw e;
707-
}
708-
}
709-
}
710599
}

0 commit comments

Comments
 (0)