Skip to content

Commit

Permalink
feat: Generate code via AI
Browse files Browse the repository at this point in the history
  • Loading branch information
@kgilpin
kgilpin committed Sep 27, 2023
1 parent e979a89 commit 0dc9754
Show file tree
Hide file tree
Showing 8 changed files with 384 additions and 219 deletions.
1 change: 1 addition & 0 deletions extensions.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
[{"identifier":{"id":"undefined_publisher.undefined"},"location":{"$mid":1,"fsPath":"/Users/kgilpin/source/appland/vscode-appland/extern","path":"/Users/kgilpin/source/appland/vscode-appland/extern","scheme":"file"},"relativeLocation":"extern"},{"identifier":{"id":"appland.appmap","uuid":"41d86b02-68d3-4049-9422-95da6d11cc2e"},"version":"0.90.1","location":{"$mid":1,"fsPath":"/Users/kgilpin/source/appland/vscode-appland/out","path":"/Users/kgilpin/source/appland/vscode-appland/out","scheme":"file"},"relativeLocation":"out","metadata":{"id":"41d86b02-68d3-4049-9422-95da6d11cc2e","publisherDisplayName":"AppLand","publisherId":"f7f1004e-6038-49cd-a096-4e618fe53f77","isPreReleaseVersion":false}}]
4 changes: 4 additions & 0 deletions package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,10 @@
"command": "appmap.ask",
"title": "AppMap: Ask a Question About the Current Code Selection"
},
{
"command": "appmap.generate",
"title": "AppMap: Generate Code Based on the Current Code Selection"
},
{
"command": "appmap.touchOutOfDateTestFiles",
"title": "AppMap: Touch Out-of-Date Test Files"
Expand Down
169 changes: 57 additions & 112 deletions src/commands/ask.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ import * as vscode from 'vscode';
import LineInfoIndex, { LineInfo } from '../services/lineInfoIndex';
import assert from 'assert';
import AppMapCollection from '../services/appmapCollection';
import ClassMapIndex from '../services/classMapIndex';
import { promptForAppMap } from '../lib/promptForAppMap';
import { debug } from 'console';
import { readFile, writeFile } from 'fs/promises';
Expand All @@ -15,10 +14,12 @@ import {
buildAppMap,
} from '@appland/models';
import buildOpenAIApi from '../lib/buildOpenAIApi';
import ask from '../lib/ask';
import { isAbsolute, join } from 'path';
import { fileExists } from '../util';
import { join } from 'path';
import { FormatType, Specification, buildDiagram, format } from '@appland/sequence-diagram';
import { randomUUID } from 'crypto';
import { Completion, Question } from '../lib/Ask';
import selectedCode from '../lib/ask/selectedCode';
import contextAppMap from '../lib/ask/contextAppMap';

class LineInfoQuickPickItem implements vscode.QuickPickItem {
constructor(public lineInfo: LineInfo) {}
Expand All @@ -34,42 +35,45 @@ class LineInfoQuickPickItem implements vscode.QuickPickItem {
export default function register(
context: vscode.ExtensionContext,
lineInfoIndex: LineInfoIndex,
classMapIndex: ClassMapIndex,
appmapCollection: AppMapCollection
): void {
context.subscriptions.push(
vscode.commands.registerCommand('appmap.ask', async () => {
const { activeTextEditor } = vscode.window;
if (!activeTextEditor) return;

const { selection } = activeTextEditor;
const selection = selectedCode(activeTextEditor);
if (!selection) return;

const startLine = selection.anchor.line;
const endLine = selection.active.line;
let selectedCode = activeTextEditor.document.getText(selection).trim();
if (selectedCode === '')
selectedCode = activeTextEditor.document.lineAt(startLine).text.trim();

const documentUri = activeTextEditor.document.uri;
const workspaceFolder = vscode.workspace.getWorkspaceFolder(documentUri);
if (!workspaceFolder) return;

const lineInfo = (await lineInfoIndex.lineInfo(documentUri)).filter(
(line) => line.codeObjects
);
const workspaceFolder = vscode.workspace.getWorkspaceFolder(documentUri);
if (!workspaceFolder) return;
if (lineInfo.length === 0) {
vscode.window.showErrorMessage(
`I couldn't find any AppMaps related to ${documentUri.path}`
);
return;
}

const lineCodePrompt =
selectedCode.split('\n').length === 0
? selectedCode
: [selectedCode.split('\n')[0], '...'].join('');
selection.text.split('\n').length === 0
? selection.text
: [selection.text.split('\n')[0], '...'].join('');
const question = await vscode.window.showInputBox({
placeHolder: `Ask a question about: ${lineCodePrompt}`,
value: 'How does this code work?',
});
if (!question) return;

const distanceFromSelection = (item: LineInfoQuickPickItem): number =>
Math.min(Math.abs(item.lineInfo.line - startLine), Math.abs(item.lineInfo.line - endLine));
Math.min(
Math.abs(item.lineInfo.line - selection.startLine),
Math.abs(item.lineInfo.line - selection.endLine)
);

const lineInfoItems = lineInfo.map((info) => new LineInfoQuickPickItem(info));
const lineDistance = lineInfoItems.map(distanceFromSelection);
Expand Down Expand Up @@ -114,16 +118,14 @@ export default function register(
appMapFileName = selectedAppMap.fsPath;
}

const q = new Question(selection.text, question);

const appmapUri = vscode.Uri.file(appMapFileName);
let codeObject: CodeObject | undefined;
let scopeCodeObjects: CodeObject[];
let returnValues: string[];
let functions: string[];
let sequenceDiagram: string;
{
const data = await readFile(appmapUri.fsPath, 'utf-8');
const appmap = buildAppMap().source(data).build();

let codeObject: CodeObject | undefined;
appmap.classMap.visit((co) => {
if (co.fqid === codeObjectEntry.fqid) codeObject = co;
});
Expand All @@ -133,14 +135,15 @@ export default function register(
);
return;
}
q.codeObject = codeObject;

const { events } = appmap;
let filterAppMap: AppMap;
{
const codeObjectEvents = events.filter(
(event) => event.codeObject.fqid === codeObjectEntry.fqid
);
returnValues = [
q.returnValues = [
...new Set(
codeObjectEvents
.map((e) => {
Expand All @@ -151,83 +154,14 @@ export default function register(
.filter(Boolean)
),
] as string[];
const roots: Event[] = [];
const filterEvents: Event[] = [...codeObjectEvents];
for (const event of codeObjectEvents) {
for (const ancestor of new EventNavigator(event).ancestors()) {
filterEvents.push(ancestor.event);
if (ancestor.event.httpServerRequest) {
roots.push(ancestor.event);
break;
}
if (!ancestor.event.parent) roots.push(ancestor.event);
}
for (const child of event.children) {
filterEvents.push(child);
for (const child2 of child.children) {
filterEvents.push(child2);
for (const child3 of child2.children) {
filterEvents.push(child3);
for (const child4 of child3.children) filterEvents.push(child4);
}
}
}
}
for (const root of roots) {
for (const child of root.children) {
filterEvents.push(child);
for (const child2 of child.children) {
filterEvents.push(child2);
for (const child3 of child2.children) {
filterEvents.push(child3);
for (const child4 of child3.children) filterEvents.push(child4);
}
}
}
}

scopeCodeObjects = [...new Set(roots.map((e) => e.codeObject))];
// const relatedCodeObjects = [...new Set(filterEvents.map((e) => e.codeObject))];
// KEG: I want to avoid sending too much source code to the AI, as it places undue priority
// on the source code rather than the code flow.
functions = [];
for (const co of [codeObject] /* relatedCodeObjects */) {
const codeObjectEntry = await classMapIndex.lookupCodeObject(co.fqid);
if (!codeObjectEntry) continue;
if (!codeObjectEntry.path || !codeObjectEntry?.lineNo) continue;

let { path } = codeObjectEntry;
const { lineNo } = codeObjectEntry;
if (!isAbsolute(path)) path = join(codeObjectEntry.folder.uri.fsPath, path);
if (!(await fileExists(path))) continue;
if (!path.startsWith(codeObjectEntry.folder.uri.fsPath)) continue;
if (
path.slice(codeObjectEntry.folder.uri.fsPath.length).startsWith('/node_modules/') ||
path.slice(codeObjectEntry.folder.uri.fsPath.length).startsWith('/vendor/')
)
continue;

const code = (await readFile(path, 'utf-8')).split('\n');
// Sort lineInfo by distance from codeObject
const distanceFromCodeObject = (lineInfo: LineInfo) => lineInfo.line - lineNo;
const nextLineInfo = lineInfo
.sort((a, b) => distanceFromCodeObject(a) - distanceFromCodeObject(b))
.filter((li) => distanceFromCodeObject(li) > 0)
.sort()[0];
const lastLine = nextLineInfo ? nextLineInfo.line : codeObjectEntry.lineNo + 10;
const codeLines = code.slice(codeObjectEntry.lineNo - 1, lastLine - 1);
functions.push([codeObjectEntry.fqid, ...codeLines].join('\n'));
}
filterAppMap = contextAppMap(appmap, codeObjectEvents);

const eventIds = new Set(filterEvents.map((e) => e.id));
filterAppMap = buildAppMap({
events: events.filter((e) => eventIds.has(e.callEvent.id)),
classMap: appmap.classMap.roots.map((c) => ({ ...c.data })),
metadata: appmap.metadata,
}).build();
q.scopeCodeObjects = [...new Set(appmap.rootEvents().map((e) => e.codeObject))];
}

{
const filterAppMapData = JSON.stringify(filterAppMap, null, 2);
let appmapName = question.replace(/[^a-zA-Z0-9]/g, '-');
let appmapName = ['ask_appmap', randomUUID()].join('-'); // question.replace(/[^a-zA-Z0-9]/g, '-'); KEG: Can result in a too-long file name.
if (appmapName.endsWith('-')) appmapName = appmapName.slice(0, -1);
const filePath = join(workspaceFolder?.uri.fsPath, `${appmapName}.appmap.json`);
await writeFile(filePath, filterAppMapData);
Expand All @@ -236,7 +170,6 @@ export default function register(
}

const specification = Specification.build(filterAppMap, { loops: true });

assert(appmapUri);
let sequenceDiagramAppMap: AppMap;
{
Expand All @@ -246,37 +179,49 @@ export default function register(
sequenceDiagramAppMap = sequenceDiagramFilter.filter(filterAppMap, []);
}
const diagram = buildDiagram(appmapUri.fsPath, sequenceDiagramAppMap, specification);
sequenceDiagram = format(FormatType.PlantUML, diagram, appmapUri.fsPath).diagram;
q.sequenceDiagram = format(FormatType.PlantUML, diagram, appmapUri.fsPath).diagram;
}

const openAI = await buildOpenAIApi(context);
if (!openAI) return;

let completion: Completion | undefined;
await vscode.window.withProgress(
{
title: `Thinking about your question...`,
location: vscode.ProgressLocation.Notification,
},
async () => {
assert(codeObject);

try {
await ask(
question,
selectedCode,
codeObject,
returnValues,
scopeCodeObjects,
functions,
sequenceDiagram,
openAI
);
completion = await q.complete(openAI);
} catch (e) {
debug((e as any).toString());
vscode.window.showErrorMessage(`Unable to process your question: ${e}`);
}
}
);
if (!completion) return;

const fence = '```';
const responseText = [
`## ${lineCodePrompt}`,
`### You asked`,
`${fence}${question}${fence}`,
`### AI response:`,
completion.response,
`## Prompt
<details>
<summary>Click to expand</summary>
${fence}
${completion.prompt}
${fence}
</details>`,
].join('\n\n');
const newDocument = await vscode.workspace.openTextDocument({
content: responseText,
});
await vscode.window.showTextDocument(newDocument);
})
);
}
Loading

0 comments on commit 0dc9754

Please sign in to comment.