Merge branch 'main' into robertbrignull/remove_selected_model

This commit is contained in:
Robert
2023-10-18 11:43:12 +01:00
12 changed files with 205 additions and 35 deletions

View File

@@ -115,21 +115,35 @@ async function extractSourceMap() {
}
if (stacktrace.includes("at")) {
const rawSourceMaps = new Map<string, RawSourceMap>();
const rawSourceMaps = new Map<string, RawSourceMap | null>();
const mappedStacktrace = await replaceAsync(
stacktrace,
stackLineRegex,
async (match, name, file, line, column) => {
if (!rawSourceMaps.has(file)) {
const rawSourceMap: RawSourceMap = await readJSON(
resolve(sourceMapsDirectory, `${basename(file)}.map`),
);
rawSourceMaps.set(file, rawSourceMap);
try {
const rawSourceMap: RawSourceMap = await readJSON(
resolve(sourceMapsDirectory, `${basename(file)}.map`),
);
rawSourceMaps.set(file, rawSourceMap);
} catch (e: unknown) {
// If the file is not found, we will not decode it and not try reading this source map again
if (e instanceof Error && "code" in e && e.code === "ENOENT") {
rawSourceMaps.set(file, null);
} else {
throw e;
}
}
}
const sourceMap = rawSourceMaps.get(file);
if (!sourceMap) {
return match;
}
const originalPosition = await SourceMapConsumer.with(
rawSourceMaps.get(file) as RawSourceMap,
sourceMap,
null,
async function (consumer) {
return consumer.originalPositionFor({

View File

@@ -702,6 +702,10 @@ export function showQueriesPanel(): boolean {
const MODEL_SETTING = new Setting("model", ROOT_SETTING);
const FLOW_GENERATION = new Setting("flowGeneration", MODEL_SETTING);
const LLM_GENERATION = new Setting("llmGeneration", MODEL_SETTING);
const LLM_GENERATION_BATCH_SIZE = new Setting(
"llmGenerationBatchSize",
MODEL_SETTING,
);
const EXTENSIONS_DIRECTORY = new Setting("extensionsDirectory", MODEL_SETTING);
const SHOW_MULTIPLE_MODELS = new Setting("showMultipleModels", MODEL_SETTING);
@@ -725,6 +729,14 @@ export class ModelConfigListener extends ConfigListener implements ModelConfig {
return !!LLM_GENERATION.getValue<boolean>();
}
/**
* Limits the number of candidates we send to the model in each request to avoid long requests.
* Note that the model may return fewer than this number of candidates.
*/
public get llmGenerationBatchSize(): number {
return LLM_GENERATION_BATCH_SIZE.getValue<number | null>() || 10;
}
public getExtensionsDirectory(languageId: string): string | undefined {
return EXTENSIONS_DIRECTORY.getValue<string>({
languageId,

View File

@@ -17,11 +17,7 @@ import { DatabaseItem } from "../databases/local-databases";
import { Mode } from "./shared/mode";
import { CancellationTokenSource } from "vscode";
import { ModelingStore } from "./modeling-store";
// Limit the number of candidates we send to the model in each request
// to avoid long requests.
// Note that the model may return fewer than this number of candidates.
const candidateBatchSize = 20;
import { ModelConfigListener } from "../config";
/**
* The auto-modeler holds state around auto-modeling jobs and allows
@@ -36,6 +32,7 @@ export class AutoModeler {
private readonly app: App,
private readonly cliServer: CodeQLCliServer,
private readonly queryRunner: QueryRunner,
private readonly modelConfig: ModelConfigListener,
private readonly modelingStore: ModelingStore,
private readonly queryStorageDir: string,
private readonly databaseItem: DatabaseItem,
@@ -109,6 +106,9 @@ export class AutoModeler {
cancellationTokenSource: CancellationTokenSource,
): Promise<void> {
void extLogger.log(`Modeling package ${packageName}`);
const candidateBatchSize = this.modelConfig.llmGenerationBatchSize;
await withProgress(async (progress) => {
// Fetch the candidates to send to the model
const allCandidateMethods = getCandidates(mode, methods, modeledMethods);

View File

@@ -89,26 +89,26 @@ export class MethodsUsageDataProvider
getTreeItem(item: MethodsUsageTreeViewItem): TreeItem {
if (isMethodTreeViewItem(item)) {
const { method } = item;
return {
label: `${item.packageName}.${item.typeName}.${item.methodName}${item.methodParameters}`,
label: `${method.packageName}.${method.typeName}.${method.methodName}${method.methodParameters}`,
collapsibleState: TreeItemCollapsibleState.Collapsed,
iconPath: this.getModelingStatusIcon(item),
iconPath: this.getModelingStatusIcon(method),
};
} else {
const method = this.getParent(item);
if (!method || !isMethodTreeViewItem(method)) {
throw new Error("Parent not found for tree item");
}
const { method, usage } = item;
return {
label: item.label,
description: `${this.relativePathWithinDatabase(item.url.uri)} [${
item.url.startLine
}, ${item.url.endLine}]`,
label: usage.label,
description: `${this.relativePathWithinDatabase(usage.url.uri)} [${
usage.url.startLine
}, ${usage.url.endLine}]`,
collapsibleState: TreeItemCollapsibleState.None,
command: {
title: "Show usage",
command: "codeQLModelEditor.jumpToMethod",
arguments: [method, item, this.databaseItem],
arguments: [method, usage, this.databaseItem],
},
};
}
@@ -146,7 +146,7 @@ export class MethodsUsageDataProvider
getChildren(item?: MethodsUsageTreeViewItem): MethodsUsageTreeViewItem[] {
if (item === undefined) {
if (this.hideModeledMethods) {
return this.sortedTreeItems.filter((api) => !api.supported);
return this.sortedTreeItems.filter((api) => !api.method.supported);
} else {
return [...this.sortedTreeItems];
}
@@ -172,21 +172,24 @@ export class MethodsUsageDataProvider
usage: Usage,
): UsageTreeViewItem | undefined {
const method = this.sortedTreeItems.find(
(m) => m.signature === methodSignature,
(m) => m.method.signature === methodSignature,
);
if (!method) {
return undefined;
}
return method.children.find((u) => usagesAreEqual(u, usage));
return method.children.find((u) => usagesAreEqual(u.usage, usage));
}
}
type MethodTreeViewItem = Method & {
type MethodTreeViewItem = {
method: Method;
children: UsageTreeViewItem[];
};
type UsageTreeViewItem = Usage & {
type UsageTreeViewItem = {
method: Method;
usage: Usage;
parent: MethodTreeViewItem;
};
@@ -195,7 +198,7 @@ export type MethodsUsageTreeViewItem = MethodTreeViewItem | UsageTreeViewItem;
function isMethodTreeViewItem(
item: MethodsUsageTreeViewItem,
): item is MethodTreeViewItem {
return "children" in item && "usages" in item;
return "children" in item && "method" in item;
}
function usagesAreEqual(u1: Usage, u2: Usage): boolean {
@@ -225,12 +228,13 @@ function sortMethodsInGroups(methods: readonly Method[], mode: Mode): Method[] {
function createTreeItems(methods: readonly Method[]): MethodTreeViewItem[] {
return methods.map((method) => {
const newMethod: MethodTreeViewItem = {
...method,
method,
children: [],
};
newMethod.children = method.usages.map((usage) => ({
...usage,
method,
usage,
// This needs to be a reference to the parent method, not a copy of it.
parent: newMethod,
}));

View File

@@ -76,6 +76,7 @@ export class ModelEditorView extends AbstractWebview<
app,
cliServer,
queryRunner,
this.modelConfig,
modelingStore,
queryStorageDir,
databaseItem,

View File

@@ -0,0 +1,16 @@
import * as React from "react";
import { Meta, StoryFn } from "@storybook/react";
import { InProgressDropdown as InProgressDropdownComponent } from "../../view/model-editor/InProgressDropdown";
export default {
title: "CodeQL Model Editor/In Progress Dropdown",
component: InProgressDropdownComponent,
} as Meta<typeof InProgressDropdownComponent>;
const Template: StoryFn<typeof InProgressDropdownComponent> = (args) => (
<InProgressDropdownComponent />
);
export const InProgressDropdown = Template.bind({});

View File

@@ -19,6 +19,7 @@ type Props = {
value: string | undefined;
options: Array<{ value: string; label: string }>;
disabled?: boolean;
className?: string;
disabledPlaceholder?: string;
onChange?: (event: ChangeEvent<HTMLSelectElement>) => void;
@@ -40,6 +41,7 @@ export function Dropdown({
options,
disabled,
disabledPlaceholder,
className,
onChange,
...props
}: Props) {
@@ -49,6 +51,7 @@ export function Dropdown({
value={disabled ? disabledValue : value}
disabled={disabled}
onChange={onChange}
className={className}
{...props}
>
{disabled ? (

View File

@@ -1,9 +1,14 @@
import * as React from "react";
import { styled } from "styled-components";
import { Dropdown } from "../common/Dropdown";
const StyledDropdown = styled(Dropdown)`
font-style: italic;
`;
export const InProgressDropdown = () => {
return (
<Dropdown
<StyledDropdown
value="Thinking..."
options={[]}
disabled={true}

View File

@@ -126,6 +126,33 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
[method, modeledMethods, onChange],
);
const removeModelClickedHandlers = useMemo(
() =>
modeledMethods.map((_, index) => () => {
const newModeledMethods = [...modeledMethods];
newModeledMethods.splice(index, 1);
onChange(method.signature, newModeledMethods);
}),
[method, modeledMethods, onChange],
);
const handleAddModelClick = useCallback(() => {
const newModeledMethod: ModeledMethod = {
type: "none",
input: "",
output: "",
kind: "",
provenance: "manual",
signature: method.signature,
packageName: method.packageName,
typeName: method.typeName,
methodName: method.methodName,
methodParameters: method.methodParameters,
};
const newModeledMethods = [...modeledMethods, newModeledMethod];
onChange(method.signature, newModeledMethods);
}, [method, modeledMethods, onChange]);
const jumpToMethod = useCallback(
() => sendJumpToMethodMessage(method),
[method],
@@ -228,6 +255,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
key={index}
appearance="icon"
aria-label="Add new model"
onClick={handleAddModelClick}
disabled={addModelButtonDisabled}
>
<Codicon name="add" />
@@ -237,6 +265,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
key={index}
appearance="icon"
aria-label="Remove model"
onClick={removeModelClickedHandlers[index]}
>
<Codicon name="trash" />
</CodiconRow>

View File

@@ -358,4 +358,84 @@ describe(MethodRow.name, () => {
expect(removeButton?.getElementsByTagName("input")[0]).toBeEnabled();
}
});
it("can add a new model", async () => {
render({
modeledMethods: [modeledMethod],
viewState: {
...viewState,
showMultipleModels: true,
},
});
onChange.mockReset();
await userEvent.click(screen.getByLabelText("Add new model"));
expect(onChange).toHaveBeenCalledTimes(1);
expect(onChange).toHaveBeenCalledWith(method.signature, [
modeledMethod,
{
type: "none",
input: "",
output: "",
kind: "",
provenance: "manual",
signature: method.signature,
packageName: method.packageName,
typeName: method.typeName,
methodName: method.methodName,
methodParameters: method.methodParameters,
},
]);
});
it("can delete the first modeled method", async () => {
render({
modeledMethods: [
{ ...modeledMethod, type: "source" },
{ ...modeledMethod, type: "sink" },
{ ...modeledMethod, type: "none" },
{ ...modeledMethod, type: "summary" },
],
viewState: {
...viewState,
showMultipleModels: true,
},
});
onChange.mockReset();
await userEvent.click(screen.getAllByLabelText("Remove model")[0]);
expect(onChange).toHaveBeenCalledTimes(1);
expect(onChange).toHaveBeenCalledWith(method.signature, [
{ ...modeledMethod, type: "sink" },
{ ...modeledMethod, type: "none" },
{ ...modeledMethod, type: "summary" },
]);
});
it("can delete a modeled method in the middle", async () => {
render({
modeledMethods: [
{ ...modeledMethod, type: "source" },
{ ...modeledMethod, type: "sink" },
{ ...modeledMethod, type: "none" },
{ ...modeledMethod, type: "summary" },
],
viewState: {
...viewState,
showMultipleModels: true,
},
});
onChange.mockReset();
await userEvent.click(screen.getAllByLabelText("Remove model")[2]);
expect(onChange).toHaveBeenCalledTimes(1);
expect(onChange).toHaveBeenCalledWith(method.signature, [
{ ...modeledMethod, type: "source" },
{ ...modeledMethod, type: "sink" },
{ ...modeledMethod, type: "summary" },
]);
});
});

View File

@@ -246,12 +246,13 @@ describe("MethodsUsageDataProvider", () => {
const usage = createUsage({});
const methodTreeItem: MethodsUsageTreeViewItem = {
...supportedMethod,
method: supportedMethod,
children: [],
};
const usageTreeItem: MethodsUsageTreeViewItem = {
...usage,
method: supportedMethod,
usage,
parent: methodTreeItem,
};
methodTreeItem.children = [usageTreeItem];
@@ -383,7 +384,9 @@ describe("MethodsUsageDataProvider", () => {
expect(
dataProvider
.getChildren()
.map((item) => (item as Method).signature),
.map(
(item) => (item as MethodsUsageTreeViewItem).method.signature,
),
).toEqual(["b.a.C.d()", "b.a.C.b()", "b.a.C.a()", "a.b.C.d()"]);
// reasoning for sort order:
// b.a.C.d() has more usages than b.a.C.b()

View File

@@ -86,7 +86,10 @@ describe("MethodsUsagePanel", () => {
await panel.revealItem(method.signature, usage);
expect(mockTreeView.reveal).toHaveBeenCalledWith(
expect.objectContaining(usage),
expect.objectContaining({
method,
usage,
}),
);
});