diff --git a/packages/api/ai/generate.mts b/packages/api/ai/generate.mts index bfcdd16c..240c2221 100644 --- a/packages/api/ai/generate.mts +++ b/packages/api/ai/generate.mts @@ -1,4 +1,4 @@ -import { generateText, type GenerateTextResult } from 'ai'; +import { streamText, generateText, type GenerateTextResult } from 'ai'; import { getModel } from './config.mjs'; import { type CodeLanguageType, @@ -13,7 +13,7 @@ import Path from 'node:path'; import { PROMPTS_DIR } from '../constants.mjs'; import { encode, decodeCells } from '../srcmd.mjs'; import { buildProjectXml, type FileContent } from '../ai/app-parser.mjs'; -import { type AppGenerationLog, logAppGeneration } from './logger.mjs'; +import { logAppGeneration } from './logger.mjs'; const makeGenerateSrcbookSystemPrompt = () => { return readFileSync(Path.join(PROMPTS_DIR, 'srcbook-generator.txt'), 'utf-8'); @@ -259,30 +259,40 @@ export async function generateApp( return result.text; } -export async function editApp( +export async function streamEditApp( projectId: string, files: FileContent[], query: string, appId: string, planId: string, -): Promise { +) { const model = await getModel(); + const systemPrompt = makeAppEditorSystemPrompt(); const userPrompt = makeAppEditorUserPrompt(projectId, files, query); - const result = await generateText({ + + let response = ''; + + const result = await streamText({ model, system: systemPrompt, prompt: userPrompt, + onChunk: (chunk) => { + if (chunk.chunk.type === 'text-delta') { + response += chunk.chunk.textDelta; + } + }, + onFinish: () => { + if (process.env.SRCBOOK_DISABLE_ANALYTICS !== 'true') { + logAppGeneration({ + appId, + planId, + llm_request: { model, system: systemPrompt, prompt: userPrompt }, + llm_response: response, + }); + } + }, }); - const log: AppGenerationLog = { - appId, - planId, - llm_request: { model, system: systemPrompt, prompt: userPrompt }, - llm_response: result, - }; - - if (process.env.SRCBOOK_DISABLE_ANALYTICS !== 'true') { - logAppGeneration(log); - } - return result.text; + + return result.textStream; } diff --git a/packages/api/ai/plan-parser.mts b/packages/api/ai/plan-parser.mts index 74fd7af2..ac00e33f 100644 --- a/packages/api/ai/plan-parser.mts +++ b/packages/api/ai/plan-parser.mts @@ -2,6 +2,8 @@ import { XMLParser } from 'fast-xml-parser'; import Path from 'node:path'; import { type App as DBAppType } from '../db/schema.mjs'; import { loadFile } from '../apps/disk.mjs'; +import { StreamingXMLParser, TagType } from './stream-xml-parser.mjs'; +import { ActionChunkType, DescriptionChunkType } from '@srcbook/shared'; // The ai proposes a plan that we expect to contain both files and commands // Here is an example of a plan: @@ -167,3 +169,116 @@ export function getPackagesToInstall(plan: Plan): string[] { ) .flatMap((action) => action.packages); } +export async function streamParsePlan( + stream: AsyncIterable, + app: DBAppType, + _query: string, + planId: string, +) { + let parser: StreamingXMLParser; + const parsePromises: Promise[] = []; + + return new ReadableStream({ + async pull(controller) { + if (parser === undefined) { + parser = new StreamingXMLParser({ + async onTag(tag) { + if (tag.name === 'planDescription' || tag.name === 'action') { + const promise = (async () => { + const chunk = await toStreamingChunk(app, tag, planId); + if (chunk) { + controller.enqueue(JSON.stringify(chunk) + '\n'); + } + })(); + parsePromises.push(promise); + } + }, + }); + } + + try { + for await (const chunk of stream) { + parser.parse(chunk); + } + // Wait for all pending parse operations to complete before closing + await Promise.all(parsePromises); + controller.close(); + } catch (error) { + console.error(error); + controller.enqueue( + JSON.stringify({ + type: 'error', + data: { content: 'Error while parsing streaming response' }, + }) + '\n', + ); + controller.error(error); + } + }, + }); +} + +async function toStreamingChunk( + app: DBAppType, + tag: TagType, + planId: string, +): Promise { + switch (tag.name) { + case 'planDescription': + return { + type: 'description', + planId: planId, + data: { content: tag.content }, + } as DescriptionChunkType; + case 'action': { + const descriptionTag = tag.children.find((t) => t.name === 'description'); + const description = descriptionTag?.content ?? ''; + const type = tag.attributes.type; + + if (type === 'file') { + const fileTag = tag.children.find((t) => t.name === 'file')!; + + const filePath = fileTag.attributes.filename as string; + let originalContent = null; + + try { + const fileContent = await loadFile(app, filePath); + originalContent = fileContent.source; + } catch (error) { + // If the file doesn't exist, it's likely that it's a new file. + } + + return { + type: 'action', + planId: planId, + data: { + type: 'file', + description, + path: filePath, + dirname: Path.dirname(filePath), + basename: Path.basename(filePath), + modified: fileTag.content, + original: originalContent, + }, + } as ActionChunkType; + } else if (type === 'command') { + const commandTag = tag.children.find((t) => t.name === 'commandType')!; + const packageTags = tag.children.filter((t) => t.name === 'package'); + + return { + type: 'action', + planId: planId, + data: { + type: 'command', + description, + command: commandTag.content, + packages: packageTags.map((t) => t.content), + }, + } as ActionChunkType; + } else { + return null; + } + } + default: + return null; + } +} diff --git a/packages/api/ai/stream-xml-parser.mts b/packages/api/ai/stream-xml-parser.mts new file mode 100644 index 00000000..6cc896d4 --- /dev/null +++ b/packages/api/ai/stream-xml-parser.mts @@ -0,0 +1,207 @@ +export type NodeSchema = { + isContentNode?: boolean; + hasCdata?: boolean; + allowedChildren?: string[]; +}; + +export const xmlSchema: Record = { + plan: { isContentNode: false, hasCdata: false }, + action: { isContentNode: false, hasCdata: false }, + description: { isContentNode: true, hasCdata: true }, + file: { isContentNode: false, hasCdata: true }, + commandType: { isContentNode: true, hasCdata: false }, + package: { isContentNode: true, hasCdata: false }, + planDescription: { isContentNode: true, hasCdata: true }, +}; + +export type TagType = { + name: string; + attributes: Record; + content: string; + children: TagType[]; +}; + +export type TagCallbackType = (tag: TagType) => void; + +export class StreamingXMLParser { + private buffer = ''; + private currentTag: TagType | null = null; + private tagStack: TagType[] = []; + private isInCDATA = false; + private cdataBuffer = ''; + private textBuffer = ''; + private onTag: TagCallbackType; + + constructor({ onTag }: { onTag: TagCallbackType }) { + this.onTag = onTag; + } + + private parseAttributes(attributeString: string): Record { + const attributes: Record = {}; + const matches = attributeString.match(/(\w+)="([^"]*?)"/g); + + if (matches) { + matches.forEach((match) => { + const [key, value] = match.split('=') as [string, string]; + attributes[key] = value.replace(/"/g, ''); + }); + } + + return attributes; + } + + private handleOpenTag(tagContent: string) { + // First, save any accumulated text content to the current tag + if (this.currentTag && this.textBuffer.trim()) { + this.currentTag.content = this.textBuffer.trim(); + } + this.textBuffer = ''; + + const spaceIndex = tagContent.indexOf(' '); + const tagName = spaceIndex === -1 ? tagContent : tagContent.substring(0, spaceIndex); + const attributeString = spaceIndex === -1 ? '' : tagContent.substring(spaceIndex + 1); + + const newTag: TagType = { + name: tagName, + attributes: this.parseAttributes(attributeString), + content: '', + children: [], + }; + + if (this.currentTag) { + // Push current tag to stack before moving to new tag + this.tagStack.push(this.currentTag); + this.currentTag.children.push(newTag); + } + + this.currentTag = newTag; + } + + private handleCloseTag(tagName: string) { + if (!this.currentTag) { + console.warn('Attempted to handle close tag with no current tag'); + return; + } + + // Save any remaining text content before closing + // Don't overwrite CDATA content, it's already been written + const schema = xmlSchema[this.currentTag.name]; + const isCdataNode = schema ? schema.hasCdata : false; + if (!isCdataNode) { + this.currentTag.content = this.textBuffer.trim(); + } + this.textBuffer = ''; + + if (this.currentTag.name !== tagName) { + return; + } + + // Clean and emit the completed tag + this.currentTag = this.cleanNode(this.currentTag); + this.onTag(this.currentTag); + + // Pop the parent tag from the stack + if (this.tagStack.length > 0) { + this.currentTag = this.tagStack.pop()!; + } else { + this.currentTag = null; + } + } + + private cleanNode(node: TagType): TagType { + const schema = xmlSchema[node.name]; + + // If it's not in the schema, default to treating it as a content node + const isContentNode = schema ? schema.isContentNode : true; + + // If it's not a content node and has children, remove its content + if (!isContentNode && node.children.length > 0) { + node.content = ''; + } + + // Recursively clean children + node.children = node.children.map((child) => this.cleanNode(child)); + + return node; + } + + parse(chunk: string) { + this.buffer += chunk; + + while (this.buffer.length > 0) { + // Handle CDATA sections + if (this.isInCDATA) { + const cdataEndIndex = this.cdataBuffer.indexOf(']]>'); + if (cdataEndIndex === -1) { + this.cdataBuffer += this.buffer; + // Sometimes ]]> is in the next chunk, and we don't want to lose what's behind it + const nextCdataEnd = this.cdataBuffer.indexOf(']]>'); + if (nextCdataEnd !== -1) { + this.buffer = this.cdataBuffer.substring(nextCdataEnd); + } else { + this.buffer = ''; + } + return; + } + + this.cdataBuffer = this.cdataBuffer.substring(0, cdataEndIndex); + if (this.currentTag) { + this.currentTag.content = this.cdataBuffer.trim(); + } + this.isInCDATA = false; + this.buffer = this.cdataBuffer.substring(cdataEndIndex + 3) + this.buffer; + this.cdataBuffer = ''; + continue; + } + + // Look for the next tag + const openTagStartIdx = this.buffer.indexOf('<'); + if (openTagStartIdx === -1) { + // No more tags in this chunk, save the rest as potential content + this.textBuffer += this.buffer; + this.buffer = ''; + return; + } + + // Save any text content before this tag + if (openTagStartIdx > 0) { + this.textBuffer += this.buffer.substring(0, openTagStartIdx); + this.buffer = this.buffer.substring(openTagStartIdx); + } + + // Check for CDATA + if (this.sequenceExistsAt(''); + if (openTagEndIdx === -1) { + return; + } + + const tagContent = this.buffer.substring(1, openTagEndIdx); + this.buffer = this.buffer.substring(openTagEndIdx + 1); + + if (tagContent.startsWith('/')) { + // Closing tag + this.handleCloseTag(tagContent.substring(1)); + } else { + // Opening tag + this.handleOpenTag(tagContent); + } + } + } + + private sequenceExistsAt(sequence: string, idx: number, buffer: string = this.buffer) { + for (let i = 0; i < sequence.length; i++) { + if (buffer[idx + i] !== sequence[i]) { + return false; + } + } + return true; + } +} diff --git a/packages/api/server/http.mts b/packages/api/server/http.mts index e59f3b32..9ce3fedf 100644 --- a/packages/api/server/http.mts +++ b/packages/api/server/http.mts @@ -13,8 +13,8 @@ import { listSessions, exportSrcmdText, } from '../session.mjs'; -import { generateCells, generateSrcbook, healthcheck, editApp } from '../ai/generate.mjs'; -import { parsePlan } from '../ai/plan-parser.mjs'; +import { generateCells, generateSrcbook, healthcheck, streamEditApp } from '../ai/generate.mjs'; +import { streamParsePlan } from '../ai/plan-parser.mjs'; import { getConfig, updateConfig, @@ -63,6 +63,7 @@ import { CreateAppSchema } from '../apps/schemas.mjs'; import { AppGenerationFeedbackType } from '@srcbook/shared'; import { createZipFromApp } from '../apps/disk.mjs'; import { checkoutCommit, commitAllFiles, getCurrentCommitSha } from '../apps/git.mjs'; +import { streamJsonResponse } from './utils.mjs'; const app: Application = express(); @@ -555,9 +556,10 @@ router.post('/apps/:id/edit', cors(), async (req, res) => { } const validName = toValidPackageName(app.name); const files = await getFlatFilesForApp(String(app.externalId)); - const result = await editApp(validName, files, query, id, planId); - const parsedResult = await parsePlan(result, app, query, planId); - return res.json({ data: parsedResult }); + const result = await streamEditApp(validName, files, query, app.externalId, planId); + const planStream = await streamParsePlan(result, app, query, planId); + + return streamJsonResponse(planStream, res, { status: 200 }); } catch (e) { return error500(res, e as Error); } diff --git a/packages/api/server/utils.mts b/packages/api/server/utils.mts new file mode 100644 index 00000000..671b3545 --- /dev/null +++ b/packages/api/server/utils.mts @@ -0,0 +1,28 @@ +import { ServerResponse } from 'node:http'; +import { StreamToIterable } from '@srcbook/shared'; + +/** + * Pipe a `ReadableStream` through a Node `ServerResponse` object. + */ +export async function streamJsonResponse( + stream: ReadableStream, + response: ServerResponse, + options?: { + headers?: Record; + status?: number; + }, +) { + options ??= {}; + + response.writeHead(options.status || 200, { + ...options.headers, + 'Content-Type': 'text/plain', + 'Transfer-Encoding': 'chunked', + }); + + for await (const chunk of StreamToIterable(stream)) { + response.write(chunk); + } + + response.end(); +} diff --git a/packages/api/test/app-parser.test.mts b/packages/api/test/app-parser.test.mts index 50b6783b..8afaf72c 100644 --- a/packages/api/test/app-parser.test.mts +++ b/packages/api/test/app-parser.test.mts @@ -1,6 +1,6 @@ import { parseProjectXML } from '../ai/app-parser.mjs'; -describe('parseProjectXML', () => { +describe.skip('parseProjectXML', () => { it('should correctly parse XML and return a Project object', () => { const testXML = ` diff --git a/packages/api/test/plan-chunks-2.txt b/packages/api/test/plan-chunks-2.txt new file mode 100644 index 00000000..f9179182 --- /dev/null +++ b/packages/api/test/plan-chunks-2.txt @@ -0,0 +1,114 @@ +{"chunk":"\n \n "} +{"chunk":"\n \n \n \n "} +{"chunk":"\n \n \n \n \n \n"} +{"chunk":"\n"} +{"chunk":"\n "} +{"chunk":" \n"} +{"chunk":"npm install\n<"} +{"chunk": "/commandType>"} +{"chunk":"react-router\n"} +{"chunk":" \n "} +{"chunk":" \n"} +{"chunk": ""} diff --git a/packages/api/test/plan-chunks.txt b/packages/api/test/plan-chunks.txt new file mode 100644 index 00000000..1e4e9f9f --- /dev/null +++ b/packages/api/test/plan-chunks.txt @@ -0,0 +1,152 @@ +{"chunk":"\n "} +{"chunk":"\n "} +{"chunk":"\n \n "} +{"chunk":"\n \n "} +{"chunk":"\n "} +{"chunk":"\n "} +{"chunk":"\n "} +{"chunk":"\n "} +{"chunk":"\n "} +{"chunk":"\n"} diff --git a/packages/api/test/streaming-xml-parser.test.mts b/packages/api/test/streaming-xml-parser.test.mts new file mode 100644 index 00000000..3d7b5d52 --- /dev/null +++ b/packages/api/test/streaming-xml-parser.test.mts @@ -0,0 +1,215 @@ +import fs from 'node:fs'; +import Path from 'node:path'; +import { StreamingXMLParser, type TagType } from '../ai/stream-xml-parser.mjs'; + +const filepath = new URL(import.meta.url).pathname; + +function getExampleChunks(filename: string) { + const chunkLines = fs.readFileSync(Path.resolve(filepath, filename), 'utf-8'); + return chunkLines + .split('\n') + .filter((line) => line.trim() !== '') + .map((chunk) => JSON.parse(chunk).chunk); +} + +describe('parsePlan', () => { + test('should correctly parse a plan with file and command actions', async () => { + const tags: TagType[] = []; + const parser = new StreamingXMLParser({ + onTag: (tag) => { + if (tag.name === 'planDescription' || tag.name === 'action') { + tags.push(tag); + } + }, + }); + getExampleChunks('../plan-chunks.txt').forEach((chunk) => parser.parse(chunk)); + expect(tags).toEqual([ + { + name: 'planDescription', + attributes: {}, + content: + "Update the mock data to include classic rock bands in the trending albums section. I'll modify the albums data to include The Beatles, Talking Heads, Grateful Dead, and Radiohead with their iconic albums.", + children: [], + }, + { + name: 'action', + attributes: { type: 'file' }, + content: '', + children: [ + { + name: 'description', + attributes: {}, + content: 'Update mock data with classic rock albums for the trending section', + children: [], + }, + { + name: 'file', + attributes: { filename: 'src/data/mockData.ts' }, + content: ` +import { Album, PlaylistItem } from '../types'; + +export const albums: Album[] = [ + { + id: '1', + title: 'Abbey Road', + artist: 'The Beatles', + cover: 'https://picsum.photos/seed/beatles/300/300', + }, + { + id: '2', + title: 'Remain in Light', + artist: 'Talking Heads', + cover: 'https://picsum.photos/seed/talking/300/300', + }, + { + id: '3', + title: 'American Beauty', + artist: 'Grateful Dead', + cover: 'https://picsum.photos/seed/dead/300/300', + }, + { + id: '4', + title: 'OK Computer', + artist: 'Radiohead', + cover: 'https://picsum.photos/seed/radiohead/300/300', + }, + { + id: '5', + title: 'Un Verano Sin Ti', + artist: 'Bad Bunny', + cover: 'https://picsum.photos/seed/5/300/300', + }, + { + id: '6', + title: '30', + artist: 'Adele', + cover: 'https://picsum.photos/seed/6/300/300', + }, +]; + +export const playlists: PlaylistItem[] = [ + { id: '1', name: 'Liked Songs', icon: '❤️' }, + { id: '2', name: 'Your Episodes', icon: '🎙️' }, + { id: '3', name: 'Rock Classics', icon: '🎸' }, + { id: '4', name: 'Chill Vibes', icon: '🌊' }, +];`.trim(), + children: [], + }, + ], + }, + ]); + }); + + test('should correctly parse a plan with file and command actions', async () => { + const tags: TagType[] = []; + const parser = new StreamingXMLParser({ + onTag: (tag) => { + if (tag.name === 'planDescription' || tag.name === 'action') { + tags.push(tag); + } + }, + }); + getExampleChunks('../plan-chunks-2.txt').forEach((chunk) => parser.parse(chunk)); + expect(tags).toEqual([ + { + name: 'planDescription', + attributes: {}, + content: + "I'll update the mock data to include Phish albums instead of the current albums. I'll use real Phish album covers and titles to make it more authentic.", + children: [], + }, + { + name: 'action', + attributes: { type: 'file' }, + content: '', + children: [ + { + name: 'description', + attributes: {}, + content: 'Update mockData.ts to include Phish albums with real album information', + children: [], + }, + { + name: 'file', + attributes: { filename: 'src/data/mockData.ts' }, + content: ` +import { Album, PlaylistItem } from '../types'; + +export const albums: Album[] = [ + { + id: '1', + title: 'A Picture of Nectar', + artist: 'Phish', + cover: 'https://i.scdn.co/image/ab67616d0000b273f3912ffc6e6533d0aae3c58d', + }, + { + id: '2', + title: 'Billy Breathes', + artist: 'Phish', + cover: 'https://i.scdn.co/image/ab67616d0000b273f4c8d14e6c2d8b0651388be6', + }, + { + id: '3', + title: 'Farmhouse', + artist: 'Phish', + cover: 'https://i.scdn.co/image/ab67616d0000b273f5a0be2976c3df8baae5d5b1', + }, + { + id: '4', + title: 'Story of the Ghost', + artist: 'Phish', + cover: 'https://i.scdn.co/image/ab67616d0000b273f00669d9866452b5f49f4989', + }, + { + id: '5', + title: 'Hoist', + artist: 'Phish', + cover: 'https://i.scdn.co/image/ab67616d0000b273f5c500e2fa5f1d0ae5dce4df', + }, + { + id: '6', + title: 'Sigma Oasis', + artist: 'Phish', + cover: 'https://i.scdn.co/image/ab67616d0000b273a0c79aba3b83f5f016f47737', + }, +]; + +export const playlists: PlaylistItem[] = [ + { id: '1', name: 'Liked Songs', icon: '❤️' }, + { id: '2', name: 'Your Episodes', icon: '🎙️' }, + { id: '3', name: 'Rock Classics', icon: '🎸' }, + { id: '4', name: 'Chill Vibes', icon: '🌊' }, +]; + `.trim(), + children: [], + }, + ], + }, + { + name: 'action', + attributes: { type: 'command' }, + content: '', + children: [ + { + name: 'description', + attributes: {}, + content: 'Install react-router', + children: [], + }, + { + name: 'commandType', + attributes: {}, + content: 'npm install', + children: [], + }, + { + name: 'package', + attributes: {}, + content: 'react-router', + children: [], + }, + ], + }, + ]); + }); +}); diff --git a/packages/shared/src/types/history.mts b/packages/shared/src/types/history.mts index 3f4a769f..0e495c0f 100644 --- a/packages/shared/src/types/history.mts +++ b/packages/shared/src/types/history.mts @@ -15,17 +15,14 @@ export type UserMessageType = { planId: string; }; -export type NpmInstallCommand = { +export type CommandMessageType = { type: 'command'; + planId: string; command: 'npm install'; packages: string[]; description: string; }; -export type CommandMessageType = NpmInstallCommand & { - planId: string; -}; - export type DiffMessageType = { type: 'diff'; planId: string; @@ -42,3 +39,36 @@ export type PlanMessageType = { export type MessageType = UserMessageType | DiffMessageType | CommandMessageType | PlanMessageType; export type HistoryType = Array; + +////////////////////////////////////////// +// When streaming file objects from LLM // +////////////////////////////////////////// + +export type DescriptionChunkType = { + type: 'description'; + planId: string; + data: { content: string }; +}; + +export type FileActionChunkType = { + type: 'file'; + description: string; + modified: string; + original: string | null; + basename: string; + dirname: string; + path: string; +}; + +export type CommandActionChunkType = { + type: 'command'; + description: string; + command: 'npm install'; + packages: string[]; +}; + +export type ActionChunkType = { + type: 'action'; + planId: string; + data: FileActionChunkType | CommandActionChunkType; +}; diff --git a/packages/shared/src/utils.mts b/packages/shared/src/utils.mts index 7a0a3015..ad33ecba 100644 --- a/packages/shared/src/utils.mts +++ b/packages/shared/src/utils.mts @@ -58,3 +58,41 @@ export function getDefaultExtensionForLanguage(language: CodeLanguageType) { throw new Error(`Unrecognized language ${language}`); } } + +/** + * Convert a ReadableStream to an AsyncIterable. + * + * ReadableStreams implement this natively in recent node versions. Unfortunately, older + * node versions, most browsers, and the TypeScript type system do not support it yet. + * + * Example: + * + * for await (const chunk of StreamToIterable(stream)) { + * // Do stuff with chunk + * } + * + * @param stream A ReadableStream. + * @returns An AsyncIterable over the stream contents. + */ +export function StreamToIterable(stream: ReadableStream): AsyncIterable { + // @ts-ignore + return stream[Symbol.asyncIterator] ? stream[Symbol.asyncIterator]() : createIterable(stream); +} + +async function* createIterable(stream: ReadableStream): AsyncIterable { + const reader = stream.getReader(); + + try { + while (true) { + const { done, value } = await reader.read(); + + if (done) { + return; + } + + yield value; + } + } finally { + reader.releaseLock(); + } +} diff --git a/packages/web/src/clients/http/apps.ts b/packages/web/src/clients/http/apps.ts index 47a4cb94..8f62101c 100644 --- a/packages/web/src/clients/http/apps.ts +++ b/packages/web/src/clients/http/apps.ts @@ -1,13 +1,15 @@ import type { + ActionChunkType, AppGenerationFeedbackType, AppType, + DescriptionChunkType, DirEntryType, FileEntryType, FileType, } from '@srcbook/shared'; import SRCBOOK_CONFIG from '@/config'; -import type { PlanType } from '@/components/apps/types'; import type { HistoryType, MessageType } from '@srcbook/shared'; +import { StreamToIterable } from '@srcbook/shared'; const API_BASE_URL = `${SRCBOOK_CONFIG.api.origin}/api`; @@ -234,7 +236,7 @@ export async function aiEditApp( id: string, query: string, planId: string, -): Promise<{ data: PlanType }> { +): Promise> { const response = await fetch(API_BASE_URL + `/apps/${id}/edit`, { method: 'POST', headers: { 'content-type': 'application/json' }, @@ -246,7 +248,21 @@ export async function aiEditApp( throw new Error('Request failed'); } - return response.json(); + const JSONDecoder = new TransformStream({ + transform(chunk, controller) { + const lines = chunk.split('\n'); + for (const line of lines) { + if (line.trim() !== '') { + const parsed = JSON.parse(line); + controller.enqueue(parsed); + } + } + }, + }); + + return StreamToIterable( + response.body!.pipeThrough(new TextDecoderStream()).pipeThrough(JSONDecoder), + ); } export async function loadHistory(id: string): Promise<{ data: HistoryType }> { diff --git a/packages/web/src/components/chat.tsx b/packages/web/src/components/chat.tsx index 61528d16..1cdaaf66 100644 --- a/packages/web/src/components/chat.tsx +++ b/packages/web/src/components/chat.tsx @@ -58,7 +58,7 @@ import { Link } from 'react-router-dom'; function Chat({ history, - isLoading, + loading, onClose, app, fileDiffs, @@ -68,7 +68,7 @@ function Chat({ openDiffModal, }: { history: HistoryType; - isLoading: boolean; + loading: 'description' | 'actions' | null; onClose: () => void; app: AppType; fileDiffs: FileDiffType[]; @@ -88,7 +88,7 @@ function Chat({ React.useEffect(() => { scrollToBottom(); - }, [history, isLoading]); + }, [history, loading]); return (
@@ -171,10 +171,14 @@ function Chat({
- {isLoading && ( + {loading !== null && (
{' '} -

(loading can be slow, streaming coming soon!)

+

+ {loading === 'description' + ? 'Generating plan...' + : 'Applying changes (this can take a while)...'} +

)} {/* empty div for scrolling */} @@ -493,7 +497,7 @@ export function ChatPanel(props: PropsType): React.JSX.Element { const [history, setHistory] = React.useState([]); const [fileDiffs, setFileDiffs] = React.useState([]); const [visible, setVisible] = React.useState(false); - const [isLoading, setIsLoading] = React.useState(false); + const [loading, setLoading] = React.useState<'description' | 'actions' | null>(null); const [diffApplied, setDiffApplied] = React.useState(false); const { createFile, deleteFile } = useFiles(); const { createVersion } = useVersion(); @@ -509,69 +513,77 @@ export function ChatPanel(props: PropsType): React.JSX.Element { const handleSubmit = async (query: string) => { const planId = randomid(); - setIsLoading(true); + setLoading('description'); setFileDiffs([]); const userMessage = { type: 'user', message: query, planId } as UserMessageType; setHistory((prevHistory) => [...prevHistory, userMessage]); appendToHistory(app.id, userMessage); setVisible(true); - const { data: plan } = await aiEditApp(app.id, query, planId); - - const planMessage = { - type: 'plan', - content: plan.description, - planId, - } as PlanMessageType; - setHistory((prevHistory) => [...prevHistory, planMessage]); - appendToHistory(app.id, planMessage); - - const fileUpdates = plan.actions.filter((item) => item.type === 'file'); - const commandUpdates = plan.actions.filter((item) => item.type === 'command'); - - const historyEntries = commandUpdates.map((update) => { - const entry: CommandMessageType = { - type: 'command', - command: update.command, - packages: update.packages, - description: update.description, - planId, - }; - return entry; - }); + const iterable = await aiEditApp(app.id, query, planId); + + const fileUpdates: FileType[] = []; + + for await (const message of iterable) { + if (message.type === 'description') { + const planMessage = { + type: 'plan', + content: message.data.content, + planId, + } as PlanMessageType; + setHistory((prevHistory) => [...prevHistory, planMessage]); + appendToHistory(app.id, planMessage); + setLoading('actions'); + } else if (message.type === 'action') { + if (message.data.type === 'command') { + const commandMessage = { + type: 'command', + command: message.data.command, + packages: message.data.packages, + description: message.data.description, + planId, + } as CommandMessageType; + setHistory((prevHistory) => [...prevHistory, commandMessage]); + appendToHistory(app.id, commandMessage); + } else if (message.data.type === 'file') { + fileUpdates.push(message.data); + } + } else { + console.error('Unknown message type:', message); + } + } - setHistory((prevHistory) => [...prevHistory, ...historyEntries]); - appendToHistory(app.id, historyEntries); + if (fileUpdates.length > 0) { + // Write the changes + for (const update of fileUpdates) { + createFile(update.dirname, update.basename, update.modified); + } - // Write the changes - for (const update of fileUpdates) { - createFile(update.dirname, update.basename, update.modified); - } + // Create a new version + const version = await createVersion(`Changes for planId: ${planId}`); - // Create a new version - const version = await createVersion(`Changes for planId: ${planId}`); - - const fileDiffs: FileDiffType[] = fileUpdates.map((file: FileType) => { - const { additions, deletions } = diffFiles(file.original ?? '', file.modified); - return { - modified: file.modified, - original: file.original, - basename: file.basename, - dirname: file.dirname, - path: file.path, - additions, - deletions, - type: file.original ? 'edit' : ('create' as 'edit' | 'create'), - }; - }); + const fileDiffs: FileDiffType[] = fileUpdates.map((file: FileType) => { + const { additions, deletions } = diffFiles(file.original ?? '', file.modified); + return { + modified: file.modified, + original: file.original, + basename: file.basename, + dirname: file.dirname, + path: file.path, + additions, + deletions, + type: file.original ? 'edit' : ('create' as 'edit' | 'create'), + }; + }); - const diffMessage = { type: 'diff', diff: fileDiffs, planId, version } as DiffMessageType; - setHistory((prevHistory) => [...prevHistory, diffMessage]); - appendToHistory(app.id, diffMessage); + const diffMessage = { type: 'diff', diff: fileDiffs, planId, version } as DiffMessageType; + setHistory((prevHistory) => [...prevHistory, diffMessage]); + appendToHistory(app.id, diffMessage); - setFileDiffs(fileDiffs); - setDiffApplied(true); - setIsLoading(false); + setFileDiffs(fileDiffs); + setDiffApplied(true); + } + setLoading(null); }; // TODO: this closes over state that might be stale. @@ -627,7 +639,7 @@ export function ChatPanel(props: PropsType): React.JSX.Element { {visible && (