Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(code-streaming): added code streaming to editor while AI is writing files #213

Merged
merged 9 commits into from
Nov 12, 2024
34 changes: 19 additions & 15 deletions app/components/chat/BaseChat.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -24,16 +24,16 @@ const EXAMPLE_PROMPTS = [
{ text: 'How do I center a div?' },
];

const providerList = [...new Set(MODEL_LIST.map((model) => model.provider))]
const providerList = [...new Set(MODEL_LIST.map((model) => model.provider))];

const ModelSelector = ({ model, setModel, provider, setProvider, modelList, providerList }) => {
return (
<div className="mb-2 flex gap-2">
<select
<select
value={provider}
onChange={(e) => {
setProvider(e.target.value);
const firstModel = [...modelList].find(m => m.provider == e.target.value);
const firstModel = [...modelList].find((m) => m.provider == e.target.value);
setModel(firstModel ? firstModel.name : '');
}}
className="flex-1 p-2 rounded-lg border border-bolt-elements-borderColor bg-bolt-elements-prompt-background text-bolt-elements-textPrimary focus:outline-none focus:ring-2 focus:ring-bolt-elements-focus transition-all"
Expand All @@ -58,11 +58,13 @@ const ModelSelector = ({ model, setModel, provider, setProvider, modelList, prov
onChange={(e) => setModel(e.target.value)}
className="flex-1 p-2 rounded-lg border border-bolt-elements-borderColor bg-bolt-elements-prompt-background text-bolt-elements-textPrimary focus:outline-none focus:ring-2 focus:ring-bolt-elements-focus transition-all"
>
{[...modelList].filter(e => e.provider == provider && e.name).map((modelOption) => (
<option key={modelOption.name} value={modelOption.name}>
{modelOption.label}
</option>
))}
{[...modelList]
.filter((e) => e.provider == provider && e.name)
.map((modelOption) => (
<option key={modelOption.name} value={modelOption.name}>
{modelOption.label}
</option>
))}
</select>
</div>
);
Expand All @@ -81,10 +83,10 @@ interface BaseChatProps {
enhancingPrompt?: boolean;
promptEnhanced?: boolean;
input?: string;
model: string;
setModel: (model: string) => void;
provider: string;
setProvider: (provider: string) => void;
model?: string;
setModel?: (model: string) => void;
provider?: string;
setProvider?: (provider: string) => void;
handleStop?: () => void;
sendMessage?: (event: React.UIEvent, messageInput?: string) => void;
handleInputChange?: (event: React.ChangeEvent<HTMLTextAreaElement>) => void;
Expand Down Expand Up @@ -144,7 +146,7 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
expires: 30, // 30 days
secure: true, // Only send over HTTPS
sameSite: 'strict', // Protect against CSRF
path: '/' // Accessible across the site
path: '/', // Accessible across the site
});
} catch (error) {
console.error('Error saving API keys to cookies:', error);
Expand Down Expand Up @@ -281,7 +283,9 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
</div>
{input.length > 3 ? (
<div className="text-xs text-bolt-elements-textTertiary">
Use <kbd className="kdb px-1.5 py-0.5 rounded bg-bolt-elements-background-depth-2">Shift</kbd> + <kbd className="kdb px-1.5 py-0.5 rounded bg-bolt-elements-background-depth-2">Return</kbd> for a new line
Use <kbd className="kdb px-1.5 py-0.5 rounded bg-bolt-elements-background-depth-2">Shift</kbd> +{' '}
<kbd className="kdb px-1.5 py-0.5 rounded bg-bolt-elements-background-depth-2">Return</kbd> for
a new line
</div>
) : null}
</div>
Expand Down Expand Up @@ -315,4 +319,4 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
</div>
);
},
);
);
4 changes: 4 additions & 0 deletions app/lib/hooks/useMessageParser.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,10 @@ const messageParser = new StreamingMessageParser({

workbenchStore.runAction(data);
},
onActionStream: (data) => {
logger.trace('onActionStream', data.action);
workbenchStore.runAction(data, true);
},
},
});

Expand Down
13 changes: 8 additions & 5 deletions app/lib/runtime/action-runner.ts
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ export class ActionRunner {
});
}

async runAction(data: ActionCallbackData) {
async runAction(data: ActionCallbackData, isStreaming: boolean = false) {
const { actionId } = data;
const action = this.actions.get()[actionId];

Expand All @@ -88,19 +88,22 @@ export class ActionRunner {
if (action.executed) {
return;
}
if (isStreaming && action.type !== 'file') {
return;
}

this.#updateAction(actionId, { ...action, ...data.action, executed: true });
this.#updateAction(actionId, { ...action, ...data.action, executed: !isStreaming });

this.#currentExecutionPromise = this.#currentExecutionPromise
.then(() => {
return this.#executeAction(actionId);
return this.#executeAction(actionId, isStreaming);
})
.catch((error) => {
console.error('Action failed:', error);
});
}

async #executeAction(actionId: string) {
async #executeAction(actionId: string, isStreaming: boolean = false) {
const action = this.actions.get()[actionId];

this.#updateAction(actionId, { status: 'running' });
Expand All @@ -121,7 +124,7 @@ export class ActionRunner {
}
}

this.#updateAction(actionId, { status: action.abortSignal.aborted ? 'aborted' : 'complete' });
this.#updateAction(actionId, { status: isStreaming ? 'running' : action.abortSignal.aborted ? 'aborted' : 'complete' });
} catch (error) {
this.#updateAction(actionId, { status: 'failed', error: 'Action failed' });
logger.error(`[${action.type}]:Action failed\n\n`, error);
Expand Down
16 changes: 16 additions & 0 deletions app/lib/runtime/message-parser.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ export interface ParserCallbacks {
onArtifactOpen?: ArtifactCallback;
onArtifactClose?: ArtifactCallback;
onActionOpen?: ActionCallback;
onActionStream?: ActionCallback;
onActionClose?: ActionCallback;
}

Expand Down Expand Up @@ -118,6 +119,21 @@ export class StreamingMessageParser {

i = closeIndex + ARTIFACT_ACTION_TAG_CLOSE.length;
} else {
if ('type' in currentAction && currentAction.type === 'file') {
let content = input.slice(i);

this._options.callbacks?.onActionStream?.({
artifactId: currentArtifact.id,
messageId,
actionId: String(state.actionId - 1),
action: {
...currentAction as FileAction,
content,
filePath: currentAction.filePath,
},

});
}
break;
}
} else {
Expand Down
33 changes: 28 additions & 5 deletions app/lib/stores/workbench.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@ import { PreviewsStore } from './previews';
import { TerminalStore } from './terminal';
import JSZip from 'jszip';
import { saveAs } from 'file-saver';
import { Octokit } from "@octokit/rest";
import { Octokit, type RestEndpointMethodTypes } from "@octokit/rest";
import * as nodePath from 'node:path';
import type { WebContainerProcess } from '@webcontainer/api';

export interface ArtifactState {
Expand Down Expand Up @@ -267,16 +268,37 @@ export class WorkbenchStore {
artifact.runner.addAction(data);
}

async runAction(data: ActionCallbackData) {
async runAction(data: ActionCallbackData, isStreaming: boolean = false) {
const { messageId } = data;

const artifact = this.#getArtifact(messageId);

if (!artifact) {
unreachable('Artifact not found');
}
if (data.action.type === 'file') {
let wc = await webcontainer
const fullPath = nodePath.join(wc.workdir, data.action.filePath);
if (this.selectedFile.value !== fullPath) {
this.setSelectedFile(fullPath);
}
if (this.currentView.value !== 'code') {
this.currentView.set('code');
}
const doc = this.#editorStore.documents.get()[fullPath];
if (!doc) {
await artifact.runner.runAction(data, isStreaming);
}

artifact.runner.runAction(data);
this.#editorStore.updateFile(fullPath, data.action.content);

if (!isStreaming) {
this.resetCurrentDocument();
await artifact.runner.runAction(data);
}
} else {
artifact.runner.runAction(data);
}
}

#getArtifact(id: string) {
Expand Down Expand Up @@ -360,9 +382,10 @@ export class WorkbenchStore {
const octokit = new Octokit({ auth: githubToken });

// Check if the repository already exists before creating it
let repo
let repo: RestEndpointMethodTypes["repos"]["get"]["response"]['data']
try {
repo = await octokit.repos.get({ owner: owner, repo: repoName });
let resp = await octokit.repos.get({ owner: owner, repo: repoName });
repo = resp.data
} catch (error) {
if (error instanceof Error && 'status' in error && error.status === 404) {
// Repository doesn't exist, so create a new one
Expand Down
2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -117,5 +117,5 @@
"resolutions": {
"@typescript-eslint/utils": "^8.0.0-alpha.30"
},
"packageManager": "pnpm@9.12.2+sha512.22721b3a11f81661ae1ec68ce1a7b879425a1ca5b991c975b074ac220b187ce56c708fe5db69f4c962c989452eee76c82877f4ee80f474cebd61ee13461b6228"
"packageManager": "pnpm@9.4.0"
}