Skip to content

Commit

Permalink
Refining AI support
Browse files Browse the repository at this point in the history
  • Loading branch information
lostintangent authored Mar 3, 2024
1 parent 962fe14 commit f26f369
Show file tree
Hide file tree
Showing 4 changed files with 380 additions and 0 deletions.
4 changes: 4 additions & 0 deletions global.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
declare module "*.txt" {
const content: string;
export default content;
}
170 changes: 170 additions & 0 deletions src/ai/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,170 @@
import {
AzureKeyCredential,
OpenAIClient,
OpenAIKeyCredential,
} from "@azure/openai";
import { ExtensionContext, commands, window } from "vscode";
import * as config from "../config";
import { EXTENSION_NAME } from "../constants";
import { SwingFile, Version, store } from "../store";
import preamble from "./preamble.txt";
import { initializeStorage, storage } from "./storage";

const userPrompt = `REQUEST:
{{REQUEST}}
RESPONSE:
`;

export async function synthesizeTemplateFiles(
prompt: string,
options: { error?: string } = {}
): Promise<SwingFile[]> {
let openai: OpenAIClient;

const apiKey = await storage.getOpenAiApiKey();
const endpointUrl = config.get("ai.endpointUrl");
if (endpointUrl) {
const credential = new AzureKeyCredential(apiKey!);
openai = new OpenAIClient(endpointUrl, credential);
} else {
const credential = new OpenAIKeyCredential(apiKey!);
openai = new OpenAIClient(credential);
}

const messages = [{ role: "system", content: preamble }];

prompt = userPrompt.replace("{{REQUEST}}", prompt);

let previousVersion: Version | undefined;
if (store.history && store.history.length > 0) {
previousVersion = store.history[store.history.length - 1];
const content = previousVersion.files
.map((e) => `<<—[${e.filename}]\n${e.content}\n—>>`)
.join("\n\n");

messages.push(
{ role: "user", content: previousVersion.prompt },
{
role: "assistant",
content,
}
);

if (options.error) {
const errorPrompt = `An error occured in the code you previously provided. Could you return an updated version of the code that fixes it? You don't need to apologize or return any prose. Simply look at the error message, and reply with the updated code that includes a fix.
ERROR:
${options.error}
RESPONSE:
`;

messages.push({
role: "user",
content: errorPrompt,
});
} else {
const editPrompt = `Here's an updated version of my previous request. Detect the edits I made, modify your previous response with the neccessary code changes, and then provide the full code again, with those modifications made. You only need to reply with files that have changed. But when changing a file, you should return the entire contents of that new file. However, you can ignore any files that haven't changed, and you don't need to apologize or return any prose, or code comments indicating that no changes were made.
${prompt}`;

messages.push({
role: "user",
content: editPrompt,
});
}
} else {
messages.push({
role: "user",
content: prompt,
});
}

console.log("CS Request: %o", messages);

const model = config.get("ai.model");
const chatCompletion = await openai.getChatCompletions(
model,
// @ts-ignore
messages
);

let response = chatCompletion.choices[0].message!.content!;

// Despite asking it not to, the model will sometimes still add
// prose to the beginning of the response. We need to remove it.
const fileStart = response.indexOf("<<—[");
if (fileStart !== 0) {
response = response.slice(fileStart);
}

console.log("CS Response: %o", response);

const files = response
.split("—>>")
.filter((e) => e !== "")
.map((e) => {
e = e.trim();
const p = e.split("]\n");
return { filename: p[0].replace("<<—[", ""), content: p[1] };
})!;

// Merge the contents of files that have the same name.
const mergedFiles: SwingFile[] = [];
files.forEach((e) => {
const existing = mergedFiles.find((f) => f.filename === e.filename);
if (existing) {
existing.content += "\n\n" + e.content;
} else {
mergedFiles.push(e);
}
});

console.log("CS Files: %o", files);

// If the model generated a component, then we need to remove any script
// files that it might have also generated. Despite asking it not to!
if (files.some((e) => e.filename.startsWith("App."))) {
files.splice(files.findIndex((e) => e.filename.startsWith("script.")));
}

// Find any files in the previous files that aren't in the new files
// and add them to the new files.
if (previousVersion) {
previousVersion.files.forEach((e) => {
if (!files.some((f) => f.filename === e.filename)) {
// @ts-ignore
files.push(e);
}
});
}

store.history!.push({ prompt, files });

return files;
}

export function registerAiModule(context: ExtensionContext) {
context.subscriptions.push(
commands.registerCommand(`${EXTENSION_NAME}.setOpenAiApiKey`, async () => {
const key = await window.showInputBox({
prompt: "Enter your OpenAI API key",
placeHolder: "",
});
if (!key) return;
await storage.setOpenAiApiKey(key);
})
);

context.subscriptions.push(
commands.registerCommand(
`${EXTENSION_NAME}.clearOpenAiApiKey`,
async () => {
await storage.deleteOpenAiApiKey();
}
)
);

initializeStorage(context);
}
174 changes: 174 additions & 0 deletions src/ai/preamble.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,174 @@
You are a web coding playground that allows users to generate runnable code snippets, using a combination of HTML, JavaScript, and CSS, as well as the component style of popular web frameworks (React, Vue, Svelte, etc.).

General rules:

* When fulfilling a request for a playground, you should separate out the HTML, JavaScript, and CSS code into files called: index.html, script.js, and style.css.
* You only generate code, and offer no other description or hints to the user.
* If the user’s request doesn’t require HTML, JavaScript, or CSS to fulfill, then omit the respective file for it.
* If the only contents of a file are code comments, then omit that file from the response.
* If the user asks about a CLI or Go, then generate a file called App.go, and populate it with the Go code needed to satisfy the request

When generating HTML, follow these rules:

* Don't include the <html>, <head>, or <body> tags, as these will be automatically added by the playground's runtime environment.
* Don't include <script> or <style> tags for either script.js or style.css, as these will be automatically added to the appropriate files.

When generating JavaScript, follow these rules:

* Make sure to import any libraries you need using just the name of the library (e.g. "import * as react from "react"). The playground's runtime environment will resolve the modules correctly, and therefore, don't try to generate a URL for a CDN.
* If you use APIs from a library, make sure to reference them from the imported library name (e.g. import * as <libraryNameCamelCase> from "<library>";).
* Don't attempt to import a specific version, or file from the library. Just import it by name, and the playground will automatically use the latest version (e.g. "d3", and not "d3@7/d3.js").
* When importing a library, always use the "* as <foo>" syntax, as opposed to trying to use the default export, or a list of named exports.
* Don't import libraries that aren't actually used.
* When using React, name the file script.jsx as opposed to script.js.
* If the user requests TypeScript, then name the file script.ts, opposed to script.js. And if you're using TypeScript + React, then name the file scirpt.tsx.
* Don't add an import for the style.css CSS file, since the runtime environment will do that automatically.
* If the user asks about Svelte, then create a file called App.svelte (that includes the requested component), and don't include the script.js file.
* If the user asks about React Native, then simply create a script.js file, and don't include the index.html or style.css files.

Here are some examples of how to format your response...

---

REQUEST:
Simple hello world app, with red text, and a pop up message that says hi to the user

RESPONSE:
<<—[index.html]
<div>Hello world</div>
—>>

<<—[script.js]
alert(“hi”);
—>>

<<—[style.css]
body {
color: red;
}
—>>

---

REQUEST:
React sample that shows how to use state management

RESPONSE:
<<—[index.html]
<div id="root"></div>
—>>

<<—[script.jsx]
import * as React from "react";
import * as ReactDOM from "react-dom";

function App() {
const [name, setName] = React.useState("World");

const handleChange = (event) => {
setName(event.target.value);
};

return (
<div>
<h1>Hello {name}</h1>
<input type="text" value={name} onChange={handleChange} />
</div>
);
}

ReactDOM.render(<App />, document.getElementById("root"));
—>>

---

REQUEST:
A button labelled "Yay!" that when clicked, displays a confetti effect

RESPONSE:
<<—[index.html]
<button>Yay!</button>
—>>

<<—[script.js]
import * as confetti from "canvas-confetti";
document.querySelector("button").addEventListener("click", () => {
confetti.confetti();
});
—>>

---

REQUEST:
Hello world react app

RESPONSE:
<<—[index.html]
<div id="root"></div>
—>>

<<—[script.jsx]
import * as React from "react";
import * as ReactDOM from "react-dom";

function App() {
return <h1>Hello world</h1>;
}

ReactDOM.render(<App />, document.getElementById("root"));
—>>

<<—[style.css]
h1 {
background-color: blue;
}
—>>

---

REQUEST:
Svelte component that demonstrates how to use a cross fade animation

RESPONSE:
<<—[App.svelte]
<script>
import { fade } from 'svelte/transition';

let visible = true;
</script>

<button on:click={() => (visible = !visible)}>Toggle</button>

{#if visible}
<div transition:fade>
<h1>Hello</h1>
</div>
{:else}
<div transition:fade>
<h1>Goodbye</h1>
</div>
{/if}
—>>

---

REQUEST:
A hello world React Native app

RESPONSE:
<<—[script.js]
import * React from 'react';
import { View, Text } from 'react-native';

const App = () => {
return (
<View>
<Text>Hello World!!!!</Text>
</View>
);
};

export default App;
—>>

---
32 changes: 32 additions & 0 deletions src/ai/storage.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
import { commands, ExtensionContext } from "vscode";
import { EXTENSION_NAME } from "../constants";

const OPENAI_CONTEXT_KEY = `${EXTENSION_NAME}:hasOpenAiApiKey`;
const OPENAI_STORAGE_KEY = `${EXTENSION_NAME}:openAiApiKey`;

export interface IAiStorage {
deleteOpenAiApiKey(): Promise<void>;
getOpenAiApiKey(): Promise<string | undefined>;
setOpenAiApiKey(apiKey: string): Promise<void>;
}

export let storage: IAiStorage;
export async function initializeStorage(context: ExtensionContext) {
storage = {
async deleteOpenAiApiKey(): Promise<void> {
await context.secrets.delete(OPENAI_STORAGE_KEY);
await commands.executeCommand("setContext", OPENAI_CONTEXT_KEY, false);
},
async getOpenAiApiKey(): Promise<string | undefined> {
return context.secrets.get(OPENAI_STORAGE_KEY);
},
async setOpenAiApiKey(key: string): Promise<void> {
await context.secrets.store(OPENAI_STORAGE_KEY, key);
await commands.executeCommand("setContext", OPENAI_CONTEXT_KEY, true);
},
};

if (storage.getOpenAiApiKey() !== undefined) {
await commands.executeCommand("setContext", OPENAI_CONTEXT_KEY, true);
}
}

0 comments on commit f26f369

Please sign in to comment.