Skip to content

Commit

Permalink
Implement new way of running prompts
Browse files Browse the repository at this point in the history
  • Loading branch information
@ColinMcNeil
ColinMcNeil committed Aug 23, 2024
1 parent 0418979 commit 52413e5
Show file tree
Hide file tree
Showing 9 changed files with 169 additions and 108 deletions.
4 changes: 4 additions & 0 deletions package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,10 @@
{
"command": "docker.labs-ai-tools-vscode.delete-prompt",
"title": "Delete a saved prompt"
},
{
"command": "docker.labs-ai-tools-vscode.run-workspace-as-prompt",
"title": "Run workspace as a prompt"
}
]
},
Expand Down
51 changes: 10 additions & 41 deletions src/commands/manageSavedPrompts.ts
View file
Original file line number Diff line number Diff line change
@@ -1,56 +1,25 @@
import { spawnSync } from 'child_process';
import * as vscode from 'vscode';
import { showPromptPicker } from '../utils/promptPicker';
import { ctx } from '../extension';



export const savePrompt = async (prompt: string) => {
if (!prompt) {
prompt = (await showPromptPicker(true))?.id || "";
prompt = (await showPromptPicker())?.id || "";
}
if (prompt === "") {
return vscode.window.showErrorMessage("No prompt selected");
}

const promptImage = vscode.workspace.getConfiguration('docker.labs-ai-tools-vscode').get('prompt-image') as string;
const args = [
"run",
"--rm",
"-v",
"/var/run/docker.sock:/var/run/docker.sock",
"--mount",
"type=volume,source=docker-prompts,target=/prompts",
promptImage,
"register",
prompt
];
const result = spawnSync('docker', args);
if (result.error || result.status !== 0) {
vscode.window.showErrorMessage(`Error saving prompt ${prompt}: ${result.error || result.stderr.toString()}`);
}
else {
ctx.globalState.update('savedPrompts', [...ctx.globalState.get('savedPrompts', []), prompt]);
vscode.window.showInformationMessage(`Saved ${prompt}`);
vscode.commands.executeCommand("docker.labs-ai-tools-vscode.generate");
vscode.commands.executeCommand("docker.labs-ai-tools-vscode.run-prompt");
}
};

export const deletePrompt = (prompt: string) => {
const promptImage = vscode.workspace.getConfiguration('docker.labs-ai-tools-vscode').get('prompt-image') as string;
const args = [
"run",
"--rm",
"-v",
"/var/run/docker.sock:/var/run/docker.sock",
"--mount",
"type=volume,source=docker-prompts,target=/prompts",
promptImage,
"unregister",
prompt
];
const result = spawnSync('docker', args);
if (result.error || result.status !== 0) {
vscode.window.showErrorMessage(`Error deleting prompt ${prompt}: ${result.error || result.stderr.toString()}`);
}
else {
vscode.window.showInformationMessage(`Deleted ${prompt}`);
vscode.commands.executeCommand("docker.labs-ai-tools-vscode.generate");
}
};
ctx.globalState.update('savedPrompts', ctx.globalState.get('savedPrompts', []).filter(p => p !== prompt));
vscode.window.showInformationMessage(`Deleted ${prompt}`);
vscode.commands.executeCommand("docker.labs-ai-tools-vscode.run-prompt");
};
102 changes: 81 additions & 21 deletions src/commands/runPrompt.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,10 @@ import {
spawnSync,
} from "child_process";
import * as vscode from "vscode";
import { verifyHasOpenAIKey } from "../extension";
import { showPromptPicker } from "../utils/promptPicker";
import { preparePromptFile } from "../utils/promptFilename";
import { spawnPromptImage } from "../utils/promptRunner";
import { spawnPromptImage, writeKeyToVolume } from "../utils/promptRunner";
import { verifyHasOpenAIKey } from "./setOpenAIKey";

const START_DOCKER_COMMAND = {
'win32': 'Start-Process -NoNewWindow -Wait -FilePath "C:\\Program Files\\Docker\\Docker\\Docker Desktop.exe"',
Expand All @@ -15,8 +15,7 @@ const START_DOCKER_COMMAND = {

const DEFAULT_USER = "local-user";

export const runPrompt = (secrets: vscode.SecretStorage) => vscode.window.withProgress({ location: vscode.ProgressLocation.Window }, async progress => {

const checkDockerDesktop = () => {
// Coerce the error to have an exit code
type DockerSpawnError = Error & { code: number };

Expand All @@ -38,7 +37,6 @@ export const runPrompt = (secrets: vscode.SecretStorage) => vscode.window.withPr

// @ts-expect-error
} catch (e: DockerSpawnError) {

const platform = process.platform;
const actionItems = e.code !== -1 ? [(platform in START_DOCKER_COMMAND ? "Start Docker" : "Try again")] : ["Install Docker Desktop", "Try again"];
return vscode.window.showErrorMessage("Error starting Docker", { modal: true, detail: (e as DockerSpawnError).toString() }, ...actionItems).then(async (value) => {
Expand All @@ -49,17 +47,18 @@ export const runPrompt = (secrets: vscode.SecretStorage) => vscode.window.withPr
vscode.env.openExternal(vscode.Uri.parse("https://www.docker.com/products/docker-desktop"));
return;
case "Try again":
return vscode.commands.executeCommand("docker.labs-ai-tools-vscode.generate");
return 'RETRY';
}
});
}
};

progress.report({ increment: 0, message: "Starting..." });

const getWorkspaceFolder = async () => {
const workspaceFolders = vscode.workspace.workspaceFolders;

if (!workspaceFolders) {
return vscode.window.showErrorMessage("No workspace open");
await vscode.window.showErrorMessage("No workspace open");
return;
}

let workspaceFolder = workspaceFolders[0];
Expand All @@ -71,44 +70,105 @@ export const runPrompt = (secrets: vscode.SecretStorage) => vscode.window.withPr
ignoreFocusOut: true,
});
if (!option) {
return vscode.window.showErrorMessage("No workspace selected");
await vscode.window.showErrorMessage("No workspace selected");
return;
}
workspaceFolder = workspaceFolders[option.index];
}

await verifyHasOpenAIKey(secrets, true);
return workspaceFolder;
};

// When running local workspace as a prompt, we need to use a config value to determine the project path
const checkHasInputWorkspace = async () => {
const existingPath = vscode.workspace.getConfiguration('docker.labs-ai-tools-vscode').get('project-dir');
if (!existingPath) {
const resp = await vscode.window.showErrorMessage("No project path set in settings", "Set project path", "Cancel");
if (resp === "Set project path") {
const resp = await vscode.window.showOpenDialog({
canSelectFiles: false,
canSelectFolders: true,
canSelectMany: false,
title: "Select project path",
});
if (resp) {
await vscode.workspace.getConfiguration('docker.labs-ai-tools-vscode').update('project-dir', resp[0].fsPath);
vscode.window.showInformationMessage(`Project path set to ${resp[0].fsPath}. You can change this in settings.`);
return resp[0].fsPath;
}
}
return;
}
return existingPath;
};

export const runPrompt: (secrets: vscode.SecretStorage, localWorkspace?: boolean) => void = (secrets: vscode.SecretStorage, localWorkspace = false) => vscode.window.withProgress({ location: vscode.ProgressLocation.Window }, async progress => {

const result = await checkDockerDesktop();
if (result === 'RETRY') {
return runPrompt(secrets, localWorkspace);
}

progress.report({ increment: 0, message: "Starting..." });
progress.report({ increment: 5, message: "Checking for OpenAI key..." });

const hasOpenAIKey = await verifyHasOpenAIKey(secrets, true);
if (!hasOpenAIKey) {
return;
}

progress.report({ increment: 5, message: "Checking for workspace..." });
const workspaceFolder = await getWorkspaceFolder();
if (!workspaceFolder) {
return;
}

const promptOption = await showPromptPicker();

progress.report({ increment: 5, message: "Checking for prompt..." });
const promptOption = localWorkspace ? { id: `local://${workspaceFolder.uri.fsPath}` } : await showPromptPicker();
if (!promptOption) {
return;
}

let apiKey = await secrets.get("openAIKey");
if (localWorkspace) {
const projectPath = await checkHasInputWorkspace();
if (!projectPath) {
vscode.window.showErrorMessage("No project path set in settings");
return;
}
}

progress.report({ increment: 5, message: "Writing prompt output file..." });
const apiKey = await secrets.get("openAIKey");



const { editor, doc } = (await preparePromptFile(workspaceFolder, promptOption.id) || {});

if (!editor || !doc) {
return;
}

const writeToEditor = (text: string) => {
const edit = new vscode.WorkspaceEdit();
edit.insert(doc.uri, new vscode.Position(doc.lineCount, 0), text);
return vscode.workspace.applyEdit(edit);
};

progress.report({ increment: 5, message: "Detecting docker desktop token" });

try {
progress.report({ increment: 5, message: "Mounting secrets..." });
await writeKeyToVolume(apiKey!);
progress.report({ increment: 5, message: "Running..." });
await spawnPromptImage(promptOption.id, workspaceFolder.uri.fsPath, DEFAULT_USER, process.platform, (json) => {

await spawnPromptImage(promptOption.id, localWorkspace ? vscode.workspace.getConfiguration('docker.labs-ai-tools-vscode').get('workspace-path') || '' : workspaceFolder.uri.fsPath, DEFAULT_USER, process.platform, (json) => {
writeToEditor(JSON.stringify(json, null, 2));
});
await doc.save();
} catch (e: unknown) {
e = e as Error;
void vscode.window.showErrorMessage("Error running prompt");
await editor.edit((edit) => {
edit.insert(
new vscode.Position(editor.document.lineCount, 0),
'```json\n' + (e as Error).toString() + '\n```'
);
});
writeToEditor('```json\n' + (e as Error).toString() + '\n```');
return;
}
});
20 changes: 20 additions & 0 deletions src/commands/setOpenAIKey.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,4 +35,24 @@ export const setOpenAIKey = async (secrets: SecretStorage, skipQuickPick: boolea
window.showInformationMessage('OpenAI key deleted');
}

};


export const verifyHasOpenAIKey = async (secrets: SecretStorage, didRunAutomatically = false) => {
const openAIKey = await secrets.get('openAIKey');
if (!openAIKey) {
return await window.showErrorMessage('Model provider set to OpenAI, but no OpenAI API key found in secrets.', {
modal: didRunAutomatically
}, 'Set Key',).then(
async (res) => {
if (res === 'Set Key') {
await setOpenAIKey(secrets, true);
return true;
}
else {
return false;
}
});
}
return true;
};
35 changes: 8 additions & 27 deletions src/extension.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@ import { nativeClient } from './utils/lsp';
import { deletePrompt, savePrompt } from './commands/manageSavedPrompts';
import { spawnSync } from 'child_process';

export let ctx: vscode.ExtensionContext;

export const workspaceCommands = {} as {
[id: string]:
{
Expand All @@ -19,30 +21,9 @@ export const extensionId = 'docker.labs-ai-tools-vscode';

export const packageJSON = vscode.extensions.getExtension(extensionId)?.packageJSON;

const MAX_POLL = 10;

export const verifyHasOpenAIKey = async (secrets: vscode.SecretStorage, didRunAutomatically = false) => {
const openAIKey = await secrets.get('openAIKey');
if (!openAIKey) {
await vscode.window.showErrorMessage('Model provider set to OpenAI, but no OpenAI API key found in secrets.', {
modal: didRunAutomatically
}, 'Set Key', 'Use Ollama',).then(
async (res) => {
if (res === 'Set Key') {
await setOpenAIKey(secrets, true);
}
else if (res === 'Use Ollama') {
await vscode.workspace.getConfiguration('docker.labs-ai-tools-vscode').update('openai-base', 'Ollama');
}
else {
return false;
}
});
}
};

export async function activate(context: vscode.ExtensionContext) {

ctx = context;
let setOpenAIKeyCommand = vscode.commands.registerCommand('docker.labs-ai-tools-vscode.set-openai-api-key', () => {
setOpenAIKey(context.secrets);
});
Expand All @@ -51,11 +32,11 @@ export async function activate(context: vscode.ExtensionContext) {

spawnSync('docker', ['pull', "vonwig/prompts:latest"]);

let runPromptCommand = vscode.commands.registerCommand('docker.labs-ai-tools-vscode.generate', () => runPrompt(context.secrets));
let runPromptCommand = vscode.commands.registerCommand('docker.labs-ai-tools-vscode.run-prompt', () => runPrompt(context.secrets));

context.subscriptions.push(runPromptCommand);

let runBoundCommands = vscode.commands.registerCommand('docker.labs-ai-tools-vscode.run', runHotCommand);
let runBoundCommands = vscode.commands.registerCommand('docker.labs-ai-tools-vscode.run-commands', runHotCommand);

context.subscriptions.push(runBoundCommands);

Expand All @@ -67,9 +48,9 @@ export async function activate(context: vscode.ExtensionContext) {

context.subscriptions.push(deletePromptCommand);

if (vscode.workspace.getConfiguration('docker.labs-ai-tools-vscode').get('openai-base') === 'OpenAI') {
void verifyHasOpenAIKey(context.secrets);
}
let runWorkspaceAsPrompt = vscode.commands.registerCommand('docker.labs-ai-tools-vscode.run-workspace-as-prompt', () => runPrompt(context.secrets, true));

context.subscriptions.push(runWorkspaceAsPrompt);

nativeClient.onNotification("$bind/register", async (args: {
uri: string, blocks: {
Expand Down
4 changes: 2 additions & 2 deletions src/utils/lsp.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,12 @@ const dockerrunner = {
"-l",
'com.docker.lsp=true',
"-l",
'com.docker.lsp.extension=labs-ai-tools-vscode',
'com.docker.lsp.extension=labs-make-runbook',
"docker/lsp",
"--workspace",
"/docker",
"--profile",
"labs-ai-tools-vscode",
"labs-make-runbook",
],
};

Expand Down
7 changes: 3 additions & 4 deletions src/utils/promptFilename.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,10 @@ export const preparePromptFile = async (workspaceFolder: vscode.WorkspaceFolder,
}
}
catch (e) {
// File does not exist
const edit = new vscode.WorkspaceEdit();
edit.createFile(uri);
await vscode.workspace.applyEdit(edit);
}



const doc = await vscode.workspace.openTextDocument(uri);

const editor = await vscode.window.showTextDocument(doc);
Expand Down
Loading

0 comments on commit 52413e5

Please sign in to comment.