Skip to content

Commit 7588271

Browse files
chariskcharisk
committed
Move automodeling logic to AutoModeler class
1 parent daf389a commit 7588271

File tree

2 files changed

+181
-130
lines changed

2 files changed

+181
-130
lines changed

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

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
import { ExternalApiUsage } 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 { 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.model(externalApiUsages, modeledMethods, mode);
37+
}
38+
39+
private async model(
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+
const usages = await runAutoModelQueries({
67+
mode,
68+
candidateMethods,
69+
cliServer: this.cliServer,
70+
queryRunner: this.queryRunner,
71+
queryStorageDir: this.queryStorageDir,
72+
databaseItem: this.databaseItem,
73+
progress: (update) => progress({ ...update, maxStep }),
74+
});
75+
if (!usages) {
76+
return;
77+
}
78+
79+
progress({
80+
step: 1800,
81+
maxStep,
82+
message: "Creating request",
83+
});
84+
85+
const request = await createAutoModelV2Request(mode, usages);
86+
87+
progress({
88+
step: 2000,
89+
maxStep,
90+
message: "Sending request",
91+
});
92+
93+
const response = await this.callAutoModelApi(request);
94+
if (!response) {
95+
return;
96+
}
97+
98+
progress({
99+
step: 2500,
100+
maxStep,
101+
message: "Parsing response",
102+
});
103+
104+
const models = loadYaml(response.models, {
105+
filename: "auto-model.yml",
106+
});
107+
108+
const loadedMethods = loadDataExtensionYaml(models);
109+
if (!loadedMethods) {
110+
return;
111+
}
112+
113+
// Any candidate that was part of the response is a negative result
114+
// meaning that the canidate is not a sink for the kinds that the LLM is checking for.
115+
// For now we model this as a sink neutral method, however this is subject
116+
// to discussion.
117+
for (const candidate of candidateMethods) {
118+
if (!(candidate.signature in loadedMethods)) {
119+
loadedMethods[candidate.signature] = {
120+
type: "neutral",
121+
kind: "sink",
122+
input: "",
123+
output: "",
124+
provenance: "ai-generated",
125+
signature: candidate.signature,
126+
packageName: candidate.packageName,
127+
typeName: candidate.typeName,
128+
methodName: candidate.methodName,
129+
methodParameters: candidate.methodParameters,
130+
};
131+
}
132+
}
133+
134+
progress({
135+
step: 2800,
136+
maxStep,
137+
message: "Applying results",
138+
});
139+
140+
await this.addModeledMethods(loadedMethods);
141+
});
142+
}
143+
144+
private async callAutoModelApi(
145+
request: ModelRequest,
146+
): Promise<ModelResponse | null> {
147+
try {
148+
return await autoModelV2(this.app.credentials, request);
149+
} catch (e) {
150+
if (e instanceof RequestError && e.status === 429) {
151+
void showAndLogExceptionWithTelemetry(
152+
this.app.logger,
153+
this.app.telemetry,
154+
redactableError(e)`Rate limit hit, please try again soon.`,
155+
);
156+
return null;
157+
} else {
158+
throw e;
159+
}
160+
}
161+
}
162+
}

‎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)