mrva/vscode-codeql
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add ability to stop automodeling

Browse Source
This commit is contained in:
Charis Kyriakou
2023-08-09 10:48:59 +00:00
parent 4f7126296e
commit 4b0f59983a
8 changed files with 94 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
extensions/ql-vscode/src/common/interface-types.ts
View File

@@ -554,6 +554,11 @@ interface GenerateExternalApiFromLlmMessage {
modeledMethods: Record<string, ModeledMethod>;
}
interface StopGeneratingExternalApiFromLlmMessage {
t: "stopGeneratingExternalApiFromLlm";
dependencyName: string;
}
interface ModelDependencyMessage {
t: "modelDependency";
}
@@ -575,4 +580,5 @@ export type FromDataExtensionsEditorMessage =
| SaveModeledMethods
| GenerateExternalApiMessage
| GenerateExternalApiFromLlmMessage
| StopGeneratingExternalApiFromLlmMessage
| ModelDependencyMessage;

4
extensions/ql-vscode/src/data-extensions-editor/auto-model-codeml-queries.ts
View File

@@ -160,6 +160,7 @@ type AutoModelQueriesOptions = {
queryStorageDir: string;
progress: ProgressCallback;
cancellationTokenSource: CancellationTokenSource;
};
export type AutoModelQueriesResult = {
@@ -174,12 +175,11 @@ export async function runAutoModelQueries({
databaseItem,
queryStorageDir,
progress,
cancellationTokenSource,
}: AutoModelQueriesOptions): Promise<AutoModelQueriesResult | undefined> {
// maxStep for this part is 1500
const maxStep = 1500;
const cancellationTokenSource = new CancellationTokenSource();
const qlpack = await qlpackOfDatabase(cliServer, databaseItem);
// CodeQL needs to have access to the database to be able to retrieve the

53
extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts
View File

@@ -15,6 +15,7 @@ import { CodeQLCliServer } from "../codeql-cli/cli";
import { QueryRunner } from "../query-server";
import { DatabaseItem } from "../databases/local-databases";
import { Mode } from "./shared/mode";
import { CancellationTokenSource } from "vscode";
// Limit the number of candidates we send to the model in each request
// to avoid long requests.
@@ -22,6 +23,8 @@ import { Mode } from "./shared/mode";
const candidateBatchSize = 20;
export class AutoModeler {
private readonly jobs: Map<string, CancellationTokenSource>;
constructor(
private readonly app: App,
private readonly cliServer: CodeQLCliServer,
@@ -34,7 +37,9 @@ export class AutoModeler {
private readonly addModeledMethods: (
modeledMethods: Record<string, ModeledMethod>,
) => Promise<void>,
) {}
) {
this.jobs = new Map<string, CancellationTokenSource>();
}
public async startModeling(
dependency: string,
@@ -42,12 +47,38 @@ export class AutoModeler {
modeledMethods: Record<string, ModeledMethod>,
mode: Mode,
): Promise<void> {
await this.modelDependency(
dependency,
externalApiUsages,
modeledMethods,
mode,
);
if (this.jobs.has(dependency)) {
return;
}
const cancellationTokenSource = new CancellationTokenSource();
this.jobs.set(dependency, cancellationTokenSource);
try {
await this.modelDependency(
dependency,
externalApiUsages,
modeledMethods,
mode,
cancellationTokenSource,
);
} finally {
this.jobs.delete(dependency);
}
}
public async stopModeling(dependency: string): Promise<void> {
void extLogger.log(`Stopping modeling for dependency ${dependency}`);
const cancellationTokenSource = this.jobs.get(dependency);
if (cancellationTokenSource) {
cancellationTokenSource.cancel();
}
}
public async stopAllModeling(): Promise<void> {
for (const cancellationTokenSource of this.jobs.values()) {
cancellationTokenSource.cancel();
}
}
private async modelDependency(
@@ -55,6 +86,7 @@ export class AutoModeler {
externalApiUsages: ExternalApiUsage[],
modeledMethods: Record<string, ModeledMethod>,
mode: Mode,
cancellationTokenSource: CancellationTokenSource,
): Promise<void> {
void extLogger.log(`Modeling dependency ${dependency}`);
await withProgress(async (progress) => {
@@ -79,6 +111,10 @@ export class AutoModeler {
);
try {
for (let i = 0; i < batchNumber; i++) {
if (cancellationTokenSource.token.isCancellationRequested) {
break;
}
const start = i * candidateBatchSize;
const end = start + candidateBatchSize;
const candidatesToProcess = allCandidateMethods.slice(start, end);
@@ -100,6 +136,7 @@ export class AutoModeler {
mode,
progress,
maxStep,
cancellationTokenSource,
);
}
} finally {
@@ -114,6 +151,7 @@ export class AutoModeler {
mode: Mode,
progress: ProgressCallback,
maxStep: number,
cancellationTokenSource: CancellationTokenSource,
): Promise<void> {
const usages = await runAutoModelQueries({
mode,
@@ -123,6 +161,7 @@ export class AutoModeler {
queryStorageDir: this.queryStorageDir,
databaseItem: this.databaseItem,
progress: (update) => progress({ ...update, maxStep }),
cancellationTokenSource,
});
if (!usages) {
return;

3
extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts
View File

@@ -187,6 +187,9 @@ export class DataExtensionsEditorView extends AbstractWebview<
);
}
break;
case "stopGeneratingExternalApiFromLlm":
await this.autoModeler.stopModeling(msg.dependencyName);
break;
case "modelDependency":
await this.modelDependency();
break;

8
extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
View File

@@ -241,6 +241,13 @@ export function DataExtensionsEditor({
[],
);
const onStopGenerateFromLlmClick = useCallback((dependencyName: string) => {
vscode.postMessage({
t: "stopGeneratingExternalApiFromLlm",
dependencyName,
});
}, []);
const onOpenDatabaseClick = useCallback(() => {
vscode.postMessage({
t: "openDatabase",
@@ -345,6 +352,7 @@ export function DataExtensionsEditor({
onChange={onChange}
onSaveModelClick={onSaveModelClick}
onGenerateFromLlmClick={onGenerateFromLlmClick}
onStopGenerateFromLlmClick={onStopGenerateFromLlmClick}
onGenerateFromSourceClick={onGenerateFromSourceClick}
onModelDependencyClick={onModelDependencyClick}
/>

25
extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx
View File

@@ -89,6 +89,7 @@ type Props = {
externalApiUsages: ExternalApiUsage[],
modeledMethods: Record<string, ModeledMethod>,
) => void;
onStopGenerateFromLlmClick: (dependencyName: string) => void;
onGenerateFromSourceClick: () => void;
onModelDependencyClick: () => void;
};
@@ -105,6 +106,7 @@ export const LibraryRow = ({
onChange,
onSaveModelClick,
onGenerateFromLlmClick,
onStopGenerateFromLlmClick,
onGenerateFromSourceClick,
onModelDependencyClick,
}: Props) => {
@@ -127,6 +129,15 @@ export const LibraryRow = ({
[title, externalApiUsages, modeledMethods, onGenerateFromLlmClick],
);
const handleStopModelWithAI = useCallback(
async (e: React.MouseEvent) => {
onStopGenerateFromLlmClick(title);
e.stopPropagation();
e.preventDefault();
},
[title, onStopGenerateFromLlmClick],
);
const handleModelFromSource = useCallback(
async (e: React.MouseEvent) => {
onGenerateFromSourceClick();
@@ -167,6 +178,12 @@ export const LibraryRow = ({
);
}, [externalApiUsages, modifiedSignatures]);
const canStopAutoModeling = useMemo(() => {
return externalApiUsages.some((externalApiUsage) =>
inProgressSignatures.has(externalApiUsage.signature),
);
}, [externalApiUsages, inProgressSignatures]);
return (
<LibraryContainer>
<TitleContainer onClick={toggleExpanded} aria-expanded={isExpanded}>
@@ -185,12 +202,18 @@ export const LibraryRow = ({
</ModeledPercentage>
{hasUnsavedChanges ? <VSCodeTag>UNSAVED</VSCodeTag> : null}
</NameContainer>
{viewState.showLlmButton && (
{viewState.showLlmButton && !canStopAutoModeling && (
<VSCodeButton appearance="icon" onClick={handleModelWithAI}>
<Codicon name="lightbulb-autofix" label="Model with AI" />
&nbsp;Model with AI
</VSCodeButton>
)}
{viewState.showLlmButton && canStopAutoModeling && (
<VSCodeButton appearance="icon" onClick={handleStopModelWithAI}>
<Codicon name="debug-stop" label="Stop model with AI" />
&nbsp;Stop
</VSCodeButton>
)}
{viewState.mode === Mode.Application && (
<VSCodeButton appearance="icon" onClick={handleModelFromSource}>
<Codicon name="code" label="Model from source" />

3
extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx
View File

@@ -31,6 +31,7 @@ type Props = {
externalApiUsages: ExternalApiUsage[],
modeledMethods: Record<string, ModeledMethod>,
) => void;
onStopGenerateFromLlmClick: (dependencyName: string) => void;
onGenerateFromSourceClick: () => void;
onModelDependencyClick: () => void;
};
@@ -49,6 +50,7 @@ export const ModeledMethodsList = ({
onChange,
onSaveModelClick,
onGenerateFromLlmClick,
onStopGenerateFromLlmClick,
onGenerateFromSourceClick,
onModelDependencyClick,
}: Props) => {
@@ -93,6 +95,7 @@ export const ModeledMethodsList = ({
onChange={onChange}
onSaveModelClick={onSaveModelClick}
onGenerateFromLlmClick={onGenerateFromLlmClick}
onStopGenerateFromLlmClick={onStopGenerateFromLlmClick}
onGenerateFromSourceClick={onGenerateFromSourceClick}
onModelDependencyClick={onModelDependencyClick}
/>

2
extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/auto-model-codeml-queries.test.ts
View File

@@ -19,6 +19,7 @@ import { MethodSignature } from "../../../../src/data-extensions-editor/external
import { join } from "path";
import { exists, readFile } from "fs-extra";
import { load as loadYaml } from "js-yaml";
import { CancellationTokenSource } from "vscode-jsonrpc";
describe("runAutoModelQueries", () => {
const qlpack = {
@@ -142,6 +143,7 @@ describe("runAutoModelQueries", () => {
}),
queryStorageDir: "/tmp/queries",
progress: jest.fn(),
cancellationTokenSource: new CancellationTokenSource(),
};
const result = await runAutoModelQueries(options);