From eee47d9af93e5a1e0a8ae46087631fdb9e7e55ef Mon Sep 17 00:00:00 2001 From: Brian Hackett Date: Mon, 20 Jan 2025 09:49:36 -0800 Subject: [PATCH] Send simulation prompt commands from client (#8) --- app/components/chat/Chat.client.tsx | 174 ++++++++++++++++++---------- app/lib/replay/SimulationPrompt.ts | 11 +- app/routes/about.tsx | 6 +- app/routes/api.chat.ts | 35 +----- 4 files changed, 126 insertions(+), 100 deletions(-) diff --git a/app/components/chat/Chat.client.tsx b/app/components/chat/Chat.client.tsx index 6f292b59..6633c3fa 100644 --- a/app/components/chat/Chat.client.tsx +++ b/app/components/chat/Chat.client.tsx @@ -3,10 +3,10 @@ * Preventing TS checks with files presented in the video for a better presentation. */ import { useStore } from '@nanostores/react'; -import type { Message } from 'ai'; +import type { CreateMessage, Message } from 'ai'; import { useChat } from 'ai/react'; import { useAnimate } from 'framer-motion'; -import { memo, useCallback, useEffect, useRef, useState } from 'react'; +import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { cssTransition, toast, ToastContainer } from 'react-toastify'; import { useMessageParser, usePromptEnhancer, useShortcuts, useSnapScroll } from '~/lib/hooks'; import { description, useChatHistory } from '~/lib/persistence'; @@ -22,10 +22,11 @@ import { useSettings } from '~/lib/hooks/useSettings'; import { useSearchParams } from '@remix-run/react'; import { createSampler } from '~/utils/sampler'; import { saveProjectContents } from './Messages.client'; -import type { SimulationPromptClientData } from '~/lib/replay/SimulationPrompt'; -import { getIFrameSimulationData } from '~/lib/replay/Recording'; +import { getSimulationRecording, getSimulationEnhancedPrompt } from '~/lib/replay/SimulationPrompt'; +import { getIFrameSimulationData, type SimulationData } from '~/lib/replay/Recording'; import { getCurrentIFrame } from '../workbench/Preview'; import { getCurrentMouseData } from '../workbench/PointSelector'; +import { assert } from '~/lib/replay/ReplayProtocolClient'; const toastAnimation = cssTransition({ enter: 'animated fadeInRight', @@ -108,6 +109,26 @@ interface ChatProps { description?: string; } +let gNumAborts = 0; + +interface InjectedMessage { + message: Message; + previousId: string; +} + +function handleInjectMessages(baseMessages: Message[], injectedMessages: InjectedMessage[]) { + const messages = []; + for (const message of baseMessages) { + messages.push(message); + for (const injectedMessage of injectedMessages) { + if (injectedMessage.previousId === message.id) { + messages.push(injectedMessage.message); + } + } + } + return messages; +} + export const ChatImpl = memo( ({ description, initialMessages, storeMessageHistory, importChat, exportChat }: ChatProps) => { useShortcuts(); @@ -117,6 +138,8 @@ export const ChatImpl = memo( const [uploadedFiles, setUploadedFiles] = useState([]); // Move here const [imageDataList, setImageDataList] = useState([]); // Move here const [searchParams, setSearchParams] = useSearchParams(); + const [injectedMessages, setInjectedMessages] = useState([]); + const [simulationLoading, setSimulationLoading] = useState(false); const files = useStore(workbenchStore.files); const { promptId } = useSettings(); @@ -124,7 +147,7 @@ export const ChatImpl = memo( const [animationScope, animate] = useAnimate(); - const { messages, isLoading, input, handleInputChange, setInput, stop, append } = useChat({ + const { messages: baseMessages, isLoading, input, handleInputChange, setInput, stop, append } = useChat({ api: '/api/chat', body: { files, @@ -137,20 +160,14 @@ export const ChatImpl = memo( 'There was an error processing your request: ' + (error.message ? error.message : 'No details were returned'), ); }, - onFinish: (message, response) => { - const usage = response.usage; - - if (usage) { - console.log('Token usage:', usage); - - // You can now use the usage data as needed - } - - logger.debug('Finished streaming'); - }, initialMessages, initialInput: Cookies.get(PROMPT_COOKIE_KEY) || '', }); + + const messages = useMemo(() => { + return handleInjectMessages(baseMessages, injectedMessages); + }, [baseMessages, injectedMessages]); + useEffect(() => { const prompt = searchParams.get('prompt'); @@ -198,8 +215,10 @@ export const ChatImpl = memo( const abort = () => { stop(); + gNumAborts++; chatStore.setKey('aborted', true); workbenchStore.abortAllActions(); + setSimulationLoading(false); }; useEffect(() => { @@ -230,13 +249,54 @@ export const ChatImpl = memo( setChatStarted(true); }; + const createRecording = async (simulationData: SimulationData, repositoryContents: string) => { + let recordingId, message; + try { + recordingId = await getSimulationRecording(simulationData, repositoryContents); + message = `[Recording of the bug](https://app.replay.io/recording/${recordingId})\n\n`; + } catch (e) { + console.error("Error creating recording", e); + message = "Error creating recording."; + } + + const recordingMessage: Message = { + id: `create-recording-${messages.length}`, + role: 'assistant', + content: message, + }; + + return { recordingId, recordingMessage }; + }; + + const getEnhancedPrompt = async (recordingId: string, repositoryContents: string) => { + let enhancedPrompt, message; + try { + enhancedPrompt = await getSimulationEnhancedPrompt(recordingId, repositoryContents); + message = `Explanation of the bug: ${enhancedPrompt}`; + } catch (e) { + console.error("Error enhancing prompt", e); + message = "Error enhancing prompt."; + } + + const enhancedPromptMessage: Message = { + id: `enhanced-prompt-${messages.length}`, + role: 'assistant', + content: message, + }; + + return { enhancedPrompt, enhancedPromptMessage }; + } + const sendMessage = async (_event: React.UIEvent, messageInput?: string, simulation?: boolean) => { const _input = messageInput || input; + const numAbortsAtStart = gNumAborts; if (_input.length === 0 || isLoading) { return; } + setSimulationLoading(true); + /** * @note (delm) Usually saving files shouldn't take long but it may take longer if there * many unsaved files. In that case we need to block user input and show an indicator @@ -248,16 +308,31 @@ export const ChatImpl = memo( const { contentBase64 } = await workbenchStore.generateZipBase64(); - let simulationClientData: SimulationPromptClientData | undefined; + let simulationEnhancedPrompt: string | undefined; + if (simulation) { const simulationData = await getIFrameSimulationData(getCurrentIFrame()); - const mouseData = getCurrentMouseData(); + const { recordingId, recordingMessage } = await createRecording(simulationData, contentBase64); - simulationClientData = { - simulationData: simulationData, - repositoryContents: contentBase64, - mouseData: mouseData, - }; + if (numAbortsAtStart != gNumAborts) { + return; + } + + console.log("RecordingMessage", recordingMessage); + setInjectedMessages([...injectedMessages, { message: recordingMessage, previousId: messages[messages.length - 1].id }]); + + if (recordingId) { + const info = await getEnhancedPrompt(recordingId, contentBase64); + + if (numAbortsAtStart != gNumAborts) { + return; + } + + simulationEnhancedPrompt = info.enhancedPrompt; + + console.log("EnhancedPromptMessage", info.enhancedPromptMessage); + setInjectedMessages([...injectedMessages, { message: info.enhancedPromptMessage, previousId: messages[messages.length - 1].id }]); + } } const fileModifications = workbenchStore.getFileModifcations(); @@ -266,47 +341,28 @@ export const ChatImpl = memo( runAnimation(); - if (fileModifications !== undefined) { - /** - * If we have file modifications we append a new user message manually since we have to prefix - * the user input with the file modifications and we don't want the new user input to appear - * in the prompt. Using `append` is almost the same as `handleSubmit` except that we have to - * manually reset the input and we'd have to manually pass in file attachments. However, those - * aren't relevant here. - */ - append({ - role: 'user', - content: [ - { - type: 'text', - text: _input, - }, - ...imageDataList.map((imageData) => ({ - type: 'image', - image: imageData, - })), - ] as any, // Type assertion to bypass compiler check - }); + setSimulationLoading(false); + append({ + role: 'user', + content: [ + { + type: 'text', + text: _input, + }, + ...imageDataList.map((imageData) => ({ + type: 'image', + image: imageData, + })), + ] as any, // Type assertion to bypass compiler check + }, { body: { simulationEnhancedPrompt } }); + + if (fileModifications !== undefined) { /** * After sending a new message we reset all modifications since the model * should now be aware of all the changes. */ workbenchStore.resetAllFileModifications(); - } else { - append({ - role: 'user', - content: [ - { - type: 'text', - text: _input, - }, - ...imageDataList.map((imageData) => ({ - type: 'image', - image: imageData, - })), - ] as any, // Type assertion to bypass compiler check - }, { body: { simulationClientData } }); } setInput(''); @@ -355,7 +411,7 @@ export const ChatImpl = memo( input={input} showChat={showChat} chatStarted={chatStarted} - isStreaming={isLoading} + isStreaming={isLoading || simulationLoading} enhancingPrompt={enhancingPrompt} promptEnhanced={promptEnhanced} sendMessage={sendMessage} diff --git a/app/lib/replay/SimulationPrompt.ts b/app/lib/replay/SimulationPrompt.ts index b5ed02d1..ade4b46a 100644 --- a/app/lib/replay/SimulationPrompt.ts +++ b/app/lib/replay/SimulationPrompt.ts @@ -1,19 +1,10 @@ -// Core logic for using simulation data from remote recording to enhance +// Core logic for using simulation data from a remote recording to enhance // the AI developer prompt. -// Currently the simulation prompt is sent from the server. - import { type SimulationData, type MouseData } from './Recording'; import { assert, ProtocolClient, sendCommandDedicatedClient } from './ReplayProtocolClient'; import JSZip from 'jszip'; -// Data supplied by the client for a simulation prompt, separate from the chat input. -export interface SimulationPromptClientData { - simulationData: SimulationData; - repositoryContents: string; // base64 encoded zip file - mouseData?: MouseData; -} - interface RerecordGenerateParams { rerecordData: SimulationData; repositoryContents: string; diff --git a/app/routes/about.tsx b/app/routes/about.tsx index 623c73f6..71cf2eba 100644 --- a/app/routes/about.tsx +++ b/app/routes/about.tsx @@ -11,7 +11,11 @@ function AboutPage() {
{() => } -
Hello World! About Page
+
+ Nut is an open source fork of Bolt.new designed to help you more easily fix bugs + and make improvements to your app which AI developers struggle with. We want to be better + at cracking tough nuts, so to speak. +
); diff --git a/app/routes/api.chat.ts b/app/routes/api.chat.ts index 0ff041cc..812c1535 100644 --- a/app/routes/api.chat.ts +++ b/app/routes/api.chat.ts @@ -1,5 +1,4 @@ import { type ActionFunctionArgs } from '@remix-run/cloudflare'; -import { type SimulationPromptClientData, getSimulationEnhancedPrompt, getSimulationRecording } from '~/lib/replay/SimulationPrompt'; import { ChatStreamController } from '~/utils/chatStreamController'; import { assert } from '~/lib/replay/ReplayProtocolClient'; import { getStreamTextArguments, type Messages } from '~/lib/.server/llm/stream-text'; @@ -16,17 +15,17 @@ Focus specifically on fixing this bug. Do not guess about other problems. `; async function chatAction({ context, request }: ActionFunctionArgs) { - const { messages, files, promptId, simulationClientData } = await request.json<{ + const { messages, files, promptId, simulationEnhancedPrompt } = await request.json<{ messages: Messages; files: any; promptId?: string; - simulationClientData?: SimulationPromptClientData; + simulationEnhancedPrompt?: string; }>(); let finished: (v?: any) => void; context.cloudflare.ctx.waitUntil(new Promise((resolve) => finished = resolve)); - console.log("SimulationClientData", simulationClientData); + console.log("SimulationEnhancedPrompt", simulationEnhancedPrompt); try { const { system, messages: coreMessages } = await getStreamTextArguments({ @@ -47,37 +46,13 @@ async function chatAction({ context, request }: ActionFunctionArgs) { async start(controller) { const chatController = new ChatStreamController(controller); - let recordingId: string | undefined; - if (simulationClientData) { - try { - const { simulationData, repositoryContents } = simulationClientData; - recordingId = await getSimulationRecording(simulationData, repositoryContents); - chatController.writeText(`[Recording of the bug](https://app.replay.io/recording/${recordingId})\n\n`); - } catch (e) { - console.error("Error creating recording", e); - chatController.writeText("Error creating recording."); - } - } - - let enhancedPrompt: string | undefined; - if (recordingId) { - try { - assert(simulationClientData, "SimulationClientData is required"); - enhancedPrompt = await getSimulationEnhancedPrompt(recordingId, simulationClientData.repositoryContents); - chatController.writeText(`Enhanced prompt: ${enhancedPrompt}\n\n`); - } catch (e) { - console.error("Error enhancing prompt", e); - chatController.writeText("Error enhancing prompt."); - } - } - - if (enhancedPrompt) { + if (simulationEnhancedPrompt) { const lastMessage = coreMessages[coreMessages.length - 1]; assert(lastMessage.role == "user", "Last message must be a user message"); assert(lastMessage.content.length > 0, "Last message must have content"); const lastContent = lastMessage.content[0]; assert(typeof lastContent == "object" && lastContent.type == "text", "Last message content must be text"); - lastContent.text += `\n\n${EnhancedPromptPrefix}\n\n${enhancedPrompt}`; + lastContent.text += `\n\n${EnhancedPromptPrefix}\n\n${simulationEnhancedPrompt}`; } try {