vzcode
Version:
Multiplayer code editor system
514 lines (513 loc) • 15.7 kB
JavaScript
import { dateToTimestamp } from '@vizhub/viz-utils';
import { randomId } from '../randomId.js';
import { diff } from '../ot.js';
/**
* Ensures the chats object exists in the ShareDB document
*/
export const ensureChatsExist = (shareDBDoc) => {
if (!shareDBDoc.data.chats) {
const op = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {},
});
shareDBDoc.submitOp(op);
}
};
/**
* Ensures a specific chat exists in the ShareDB document
*/
export const ensureChatExists = (shareDBDoc, chatId) => {
if (!shareDBDoc.data.chats[chatId]) {
const op = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
id: chatId,
messages: [],
createdAt: dateToTimestamp(new Date()),
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(op);
}
};
/**
* Adds a user message to the chat
* Clears old messages to reflect that each prompt is a self-contained code transformation
*/
export const addUserMessage = (shareDBDoc, chatId, content) => {
const userMessage = {
id: `user-${Date.now()}`,
role: 'user',
content: content,
timestamp: dateToTimestamp(new Date()),
};
const userMessageOp = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
messages: [userMessage], // Replace old messages with just the new user message
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(userMessageOp);
return userMessage;
};
const DEBUG = false;
/**
* Updates AI status in the chat
*/
export const updateAIStatus = (shareDBDoc, chatId, status) => {
DEBUG &&
console.log(`ChatOperations: updateAIStatus called with status: "${status}" for chatId: ${chatId}`);
const op = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
aiStatus: status,
},
},
});
DEBUG &&
console.log(`ChatOperations: Submitting operation for status update:`, op);
shareDBDoc.submitOp(op);
DEBUG &&
console.log(`ChatOperations: Status update operation submitted successfully`);
};
/**
* Updates AI scratchpad content
*/
export const updateAIScratchpad = (shareDBDoc, chatId, content) => {
const op = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
aiScratchpad: content,
},
},
});
// op is `null` if there are no changes
// This can happen if the content is the same as before
// In that case, we don't need to submit an operation
if (op) {
shareDBDoc.submitOp(op);
}
};
/**
* Clears AI scratchpad and updates status
*/
export const clearAIScratchpadAndStatus = (shareDBDoc, chatId, status) => {
const op = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
aiScratchpad: undefined,
aiStatus: status,
},
},
});
shareDBDoc.submitOp(op);
};
/**
* Creates an initial empty AI message for streaming
*/
export const createAIMessage = (shareDBDoc, chatId) => {
const aiMessage = {
id: `assistant-${Date.now()}`,
role: 'assistant',
content: '',
timestamp: dateToTimestamp(new Date()),
};
const messageOp = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
messages: [
...shareDBDoc.data.chats[chatId].messages,
aiMessage,
],
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(messageOp);
return aiMessage.id;
};
/**
* Updates the content of an AI message during streaming
*/
export const updateAIMessageContent = (shareDBDoc, chatId, messageId, content) => {
const chat = shareDBDoc.data.chats[chatId];
const messageIndex = chat.messages.findIndex((msg) => msg.id === messageId);
if (messageIndex === -1) {
console.warn(`AI message with id ${messageId} not found`);
return;
}
const updatedMessages = [...chat.messages];
updatedMessages[messageIndex] = {
...updatedMessages[messageIndex],
content,
};
const messageOp = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...chat,
messages: updatedMessages,
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(messageOp);
};
/**
* Sets the AI status for a chat
*/
export const setAIStatus = (shareDBDoc, chatId, status) => {
const op = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
aiStatus: status,
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(op);
};
/**
* Finalizes an AI message by clearing temporary fields
*/
export const finalizeAIMessage = (shareDBDoc, chatId) => {
const op = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
aiScratchpad: undefined,
aiStatus: undefined,
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(op);
};
/**
* Adds diff data to the most recent AI message
*/
export const addDiffToAIMessage = (shareDBDoc, chatId, diffData, beforeCommitId) => {
const chat = shareDBDoc.data.chats[chatId];
const messages = [...chat.messages];
// Find the most recent AI message
const lastAIMessageIndex = messages.length - 1;
if (lastAIMessageIndex >= 0 &&
messages[lastAIMessageIndex].role === 'assistant') {
const newMessage = {
...messages[lastAIMessageIndex],
diffData,
...(beforeCommitId && { beforeCommitId }), // Add beforeCommitId for VizHub integration
};
// Use type assertion to extend the message with diffData
messages[lastAIMessageIndex] = newMessage;
const messageOp = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...chat,
messages,
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(messageOp);
}
};
/**
* Adds an AI response message to the chat (legacy function, kept for compatibility)
*/
export const addAIMessage = (shareDBDoc, chatId, content) => {
const aiResponse = {
id: Date.now() + 1,
role: 'assistant',
content: content || 'AI edit completed successfully.',
timestamp: dateToTimestamp(new Date()),
};
const messageOp = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
messages: [
...shareDBDoc.data.chats[chatId].messages,
aiResponse,
],
updatedAt: dateToTimestamp(new Date()),
aiStatus: undefined,
},
},
});
shareDBDoc.submitOp(messageOp);
return aiResponse;
};
/**
* Updates files in the ShareDB document
*/
export const updateFiles = (shareDBDoc, files) => {
const filesOp = diff(shareDBDoc.data, {
...shareDBDoc.data,
files,
});
DEBUG && console.log('updateFiles op:');
DEBUG && console.log(JSON.stringify(filesOp, null, 2));
shareDBDoc.submitOp(filesOp);
return filesOp;
};
/**
* Finds a file ID by searching for a matching file name
*/
export const resolveFileId = (fileName, shareDBDoc) => {
const files = shareDBDoc.data.files;
// Search through all files to find matching name
for (const [fileId, file] of Object.entries(files)) {
if (file.name === fileName) {
return fileId;
}
}
// If file doesn't exist, return null
return null;
};
/**
* Creates a new file with a random ID
*/
export const createNewFile = (shareDBDoc, fileName) => {
// Generate a new random file ID
const newFileId = randomId();
const newState = {
...shareDBDoc.data,
files: {
...shareDBDoc.data.files,
[newFileId]: {
name: fileName,
text: '',
},
},
};
const op = diff(shareDBDoc.data, newState);
shareDBDoc.submitOp(op);
return newFileId;
};
// ============================================================================
// Streaming Chat Operations
// ============================================================================
/**
* Creates a streaming AI message with events array
*/
export const createStreamingAIMessage = (shareDBDoc, chatId) => {
const aiMessage = {
id: `assistant-${Date.now()}`,
role: 'assistant',
content: '',
timestamp: dateToTimestamp(new Date()),
streamingEvents: [],
isProgressive: true,
};
const messageOp = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
messages: [
...shareDBDoc.data.chats[chatId].messages,
aiMessage,
],
updatedAt: dateToTimestamp(new Date()),
isStreaming: true,
},
},
});
shareDBDoc.submitOp(messageOp);
return aiMessage.id;
};
/**
* Adds a streaming event to the most recent AI message
*/
export const addStreamingEvent = (shareDBDoc, chatId, event) => {
DEBUG &&
console.log(`ChatOperations: Adding streaming event:`, event);
const chat = shareDBDoc.data.chats[chatId];
const messages = [...chat.messages];
const lastMessageIndex = messages.length - 1;
if (lastMessageIndex >= 0 &&
messages[lastMessageIndex].role === 'assistant') {
const lastMessage = messages[lastMessageIndex];
const updatedEvents = [
...(lastMessage.streamingEvents || []),
event,
];
messages[lastMessageIndex] = {
...lastMessage,
streamingEvents: updatedEvents,
};
const messageOp = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...chat,
messages,
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(messageOp);
}
};
/**
* Updates streaming status for a chat
*/
export const updateStreamingStatus = (shareDBDoc, chatId, status, isStreaming = true) => {
DEBUG &&
console.log(`ChatOperations: Updating streaming status: "${status}"`);
const op = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
currentStatus: status,
isStreaming,
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(op);
};
/**
* Sets the model for a chat when AI generation starts
*/
export const setChatModel = (shareDBDoc, chatId, model) => {
DEBUG &&
console.log(`ChatOperations: Setting model: "${model}" for chatId: ${chatId}`);
const op = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...shareDBDoc.data.chats[chatId],
model: model,
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(op);
};
/**
* Finalizes streaming message and clears streaming state
*/
export const finalizeStreamingMessage = (shareDBDoc, chatId) => {
DEBUG &&
console.log(`ChatOperations: Finalizing streaming message`);
const chat = shareDBDoc.data.chats[chatId];
const messages = [...chat.messages];
const lastMessageIndex = messages.length - 1;
if (lastMessageIndex >= 0 &&
messages[lastMessageIndex].role === 'assistant') {
const lastMessage = messages[lastMessageIndex];
messages[lastMessageIndex] = {
...lastMessage,
isComplete: true,
};
const messageOp = diff(shareDBDoc.data, {
...shareDBDoc.data,
chats: {
...shareDBDoc.data.chats,
[chatId]: {
...chat,
messages,
currentStatus: 'Done',
isStreaming: false,
updatedAt: dateToTimestamp(new Date()),
},
},
});
shareDBDoc.submitOp(messageOp);
}
};
// /**
// * Ensures a file exists, creating it if necessary
// */
// export const ensureFileExists = (shareDBDoc, fileName) => {
// let fileId = resolveFileId(fileName, shareDBDoc);
// if (!fileId) {
// // File doesn't exist, create it
// fileId = createNewFile(shareDBDoc, fileName);
// }
// return fileId;
// };
// /**
// * Clears the content of a file
// */
// export const clearFileContent = (
// shareDBDoc: ShareDBDoc<VizContent>,
// fileId: VizFileId,
// ) => {
// const currentFile = shareDBDoc.data.files[fileId];
// if (currentFile && currentFile.text) {
// // Clear the file content
// const newState = {
// ...shareDBDoc.data,
// files: {
// ...shareDBDoc.data.files,
// [fileId]: {
// ...currentFile,
// text: '',
// },
// },
// };
// const op = diff(shareDBDoc.data, newState);
// shareDBDoc.submitOp(op);
// }
// };
// /**
// * Appends a line to a file using OT operations
// */
// export const appendLineToFile = (
// shareDBDoc: ShareDBDoc<VizContent>,
// fileId: VizFileId,
// line: string,
// ) => {
// const currentFile = shareDBDoc.data.files[fileId];
// const currentContent = currentFile?.text || '';
// const newContent = currentContent + line + '\n';
// const newDocState = {
// ...shareDBDoc.data,
// files: {
// ...shareDBDoc.data.files,
// [fileId]: {
// ...currentFile,
// text: newContent,
// },
// },
// };
// shareDBDoc.submitOp(diff(shareDBDoc.data, newDocState));
// };