Send simulation prompt commands from client (#8)

This commit is contained in:
Brian Hackett 2025-01-20 09:49:36 -08:00 committed by GitHub
parent a5c5d3910a
commit eee47d9af9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 126 additions and 100 deletions

View File

@ -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<File[]>([]); // Move here
const [imageDataList, setImageDataList] = useState<string[]>([]); // Move here
const [searchParams, setSearchParams] = useSearchParams();
const [injectedMessages, setInjectedMessages] = useState<InjectedMessage[]>([]);
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,14 +341,8 @@ 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.
*/
setSimulationLoading(false);
append({
role: 'user',
content: [
@ -286,27 +355,14 @@ export const ChatImpl = memo(
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}

View File

@ -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;

View File

@ -11,7 +11,11 @@ function AboutPage() {
<BackgroundRays />
<Header />
<ClientOnly>{() => <Menu />}</ClientOnly>
<div>Hello World! About Page</div>
<div>
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.
</div>
</div>
</TooltipProvider>
);

View File

@ -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 {