From 8db2a37a59f2d09f35828e32bbcb4a3f19965584 Mon Sep 17 00:00:00 2001 From: Enrico Ros Date: Mon, 3 Jun 2024 09:21:58 -0700 Subject: [PATCH] Attachments: remove DBlobs when setting outputFragments, until GC comes --- .../attachment-drafts/attachment.dblobs.ts | 13 ++- .../attachment-drafts/attachment.pipeline.ts | 51 +++++++----- .../store-attachment-drafts-slice.tsx | 81 +++++++++++++++---- 3 files changed, 105 insertions(+), 40 deletions(-) diff --git a/src/common/attachment-drafts/attachment.dblobs.ts b/src/common/attachment-drafts/attachment.dblobs.ts index 7d083cfa8..a116de841 100644 --- a/src/common/attachment-drafts/attachment.dblobs.ts +++ b/src/common/attachment-drafts/attachment.dblobs.ts @@ -1,4 +1,4 @@ -import { addDBlobItem } from '~/modules/dblobs/dblobs.db'; +import { addDBlobItem, deleteDBlobItem } from '~/modules/dblobs/dblobs.db'; import { createDBlobImageItem } from '~/modules/dblobs/dblobs.types'; import { convertBase64Image, getImageDimensions, LLMImageResizeMode, resizeBase64ImageIfNeeded } from '~/common/util/imageUtils'; @@ -89,4 +89,13 @@ export async function attachmentImageToFragmentViaDBlob(mimeType: string, inputD console.error('imageAttachment: Error processing image:', error); return null; } -} \ No newline at end of file +} + +/** + * Remove the DBlob item associated with the given DMessageAttachmentFragment + */ +export async function removeDBlobItemFromAttachmentFragment(fragment: DMessageAttachmentFragment) { + if (fragment.part.pt === 'image_ref' && fragment.part.dataRef.reftype === 'dblob') { + await deleteDBlobItem(fragment.part.dataRef.dblobId); + } +} diff --git a/src/common/attachment-drafts/attachment.pipeline.ts b/src/common/attachment-drafts/attachment.pipeline.ts index a6eedadf2..273f698cc 100644 --- a/src/common/attachment-drafts/attachment.pipeline.ts +++ b/src/common/attachment-drafts/attachment.pipeline.ts @@ -7,6 +7,7 @@ import { pdfToImageDataURLs, pdfToText } from '~/common/util/pdfUtils'; import { createTextAttachmentFragment, DMessageAttachmentFragment } from '~/common/stores/chat/chat.message'; import type { AttachmentDraft, AttachmentDraftConverter, AttachmentDraftInput, AttachmentDraftSource } from './attachment.types'; +import type { AttachmentsDraftsStore } from './store-attachment-drafts-slice'; import { attachmentImageToFragmentViaDBlob } from './attachment.dblobs'; @@ -113,7 +114,7 @@ export function attachmentCreate(source: AttachmentDraftSource): AttachmentDraft * @param {Readonly} source - The source of the attachment. * @param {(changes: Partial) => void} edit - A function to edit the AttachmentDraft object. */ -export async function attachmentLoadInputAsync(source: Readonly, edit: (changes: Partial) => void) { +export async function attachmentLoadInputAsync(source: Readonly, edit: (changes: Partial>) => void) { edit({ inputLoading: true }); switch (source.media) { @@ -218,7 +219,7 @@ export async function attachmentLoadInputAsync(source: Readonly} input - The input of the AttachmentDraft object. * @param {(changes: Partial) => void} edit - A function to edit the AttachmentDraft object. */ -export function attachmentDefineConverters(sourceType: AttachmentDraftSource['media'], input: Readonly, edit: (changes: Partial) => void) { +export function attachmentDefineConverters(sourceType: AttachmentDraftSource['media'], input: Readonly, edit: (changes: Partial>) => void) { // return all the possible converters for the input const converters: AttachmentDraftConverter[] = []; @@ -281,16 +282,22 @@ export function attachmentDefineConverters(sourceType: AttachmentDraftSource['me * * @param {Readonly} attachment - The AttachmentDraft object to convert. * @param {number | null} converterIdx - The index of the selected converter. - * @param {(changes: Partial) => void} edit - A function to edit the AttachmentDraft object. + * @param edit - A function to edit the AttachmentDraft object. + * @param replaceOutputFragments - A function to replace the output fragments of the AttachmentDraft object. */ -export async function attachmentPerformConversion(attachment: Readonly, converterIdx: number | null, edit: (changes: Partial) => void) { +export async function attachmentPerformConversion( + attachment: Readonly, + converterIdx: number | null, + edit: AttachmentsDraftsStore['_editAttachment'], + replaceOutputFragments: AttachmentsDraftsStore['_replaceAttachmentOutputFragments'], +) { // set converter index converterIdx = (converterIdx !== null && converterIdx >= 0 && converterIdx < attachment.converters.length) ? converterIdx : null; - edit({ + edit(attachment.id, { converterIdx: converterIdx, - outputFragments: [], }); + replaceOutputFragments(attachment.id, []); // get converter const { source, ref, input } = attachment; @@ -298,23 +305,23 @@ export async function attachmentPerformConversion(attachment: Readonly')); + newFragments.push(createTextAttachmentFragment(input.altData!, ref || '\n')); break; // html to markdown table @@ -326,7 +333,7 @@ export async function attachmentPerformConversion(attachment: Readonly DMessageAttachmentFragment[]; - _editAttachment: (attachmentDraftId: AttachmentDraftId, update: Partial | ((attachment: AttachmentDraft) => Partial)) => void; + _editAttachment: (attachmentDraftId: AttachmentDraftId, update: Partial> | ((attachment: AttachmentDraft) => Partial>)) => void; + _replaceAttachmentOutputFragments: (attachmentDraftId: AttachmentDraftId, outputFragments: DMessageAttachmentFragment[]) => void; _getAttachment: (attachmentDraftId: AttachmentDraftId) => AttachmentDraft | undefined; } @@ -56,7 +58,7 @@ export const createAttachmentDraftsStoreSlice: StateCreator) => _editAttachment(attachmentDraftId, changes); + const editFn = (changes: Partial>) => _editAttachment(attachmentDraftId, changes); // 1.Resolve the Input await attachmentLoadInputAsync(source, editFn); @@ -75,13 +77,36 @@ export const createAttachmentDraftsStoreSlice: StateCreator -1 ? firstEnabledIndex : 0); }, - clearAttachmentsDrafts: () => _set({ - attachmentDrafts: [], - }), + clearAttachmentsDrafts: () => + _set(_state => { + // NOTE: commented because right now the attachments are not moved to a different scope + // because this function is actually used to clear the attachments when the message is sent + // TODO: do not use clearAttachments when the message is sent, figure out another way0 + // Remove the DBlob items associated with the removed fragments + // for (let draft of _state.attachmentDrafts) { + // for (let fragment of draft.outputFragments) { + // void removeDBlobItemFromAttachmentFragment(fragment); + // } + // } + return { + attachmentDrafts: [], + }; + }), removeAttachmentDraft: (attachmentDraftId: AttachmentDraftId) => _set(state => ({ - attachmentDrafts: state.attachmentDrafts.filter(attachment => attachment.id !== attachmentDraftId), + attachmentDrafts: state.attachmentDrafts.filter(attachment => { + if (attachment.id !== attachmentDraftId) + return true; + + // Remove the DBlob items associated with the removed fragments + for (let removedFragment of attachment.outputFragments) { + void removeDBlobItemFromAttachmentFragment(removedFragment); + } + + // Remove the draft + return false; + }), })), moveAttachmentDraft: (attachmentDraftId: AttachmentDraftId, delta: 1 | -1) => @@ -101,14 +126,12 @@ export const createAttachmentDraftsStoreSlice: StateCreator { - const { _getAttachment, _editAttachment } = _get(); + const { _getAttachment, _editAttachment, _replaceAttachmentOutputFragments } = _get(); const attachmentDraft = _getAttachment(attachmentDraftId); if (!attachmentDraft || attachmentDraft.converterIdx === converterIdx) return; - const editFn = (changes: Partial) => _editAttachment(attachmentDraftId, changes); - - await attachmentPerformConversion(attachmentDraft, converterIdx, editFn); + await attachmentPerformConversion(attachmentDraft, converterIdx, _editAttachment, _replaceAttachmentOutputFragments); }, takeAllFragments: (removeFragments: boolean): DMessageAttachmentFragment[] => { @@ -137,7 +160,7 @@ export const createAttachmentDraftsStoreSlice: StateCreator fragment.part.pt === 'text'); textFragments.push(...extractedTextFragments); @@ -147,12 +170,17 @@ export const createAttachmentDraftsStoreSlice: StateCreator fragment.part.pt !== 'text'); - if (remainingFragments.length || draft.outputsConverting) { + // Removal: rmeove associated DBlob items + for (let removedFragment of extractedTextFragments) { + void removeDBlobItemFromAttachmentFragment(removedFragment); + } + + // Removal: leave non-text fragments in the draft + const keptFragments = draft.outputFragments.filter(fragment => fragment.part.pt !== 'text'); + if (keptFragments.length || draft.outputsConverting) { keptDrafts.push({ ...draft, - outputFragments: remainingFragments, + outputFragments: keptFragments, }); } } @@ -166,7 +194,7 @@ export const createAttachmentDraftsStoreSlice: StateCreator | ((attachment: AttachmentDraft) => Partial)) => + _editAttachment: (attachmentDraftId: AttachmentDraftId, update: Partial> | ((attachment: AttachmentDraft) => Partial>)) => _set(state => ({ attachmentDrafts: state.attachmentDrafts.map((attachmentDraft: AttachmentDraft): AttachmentDraft => attachmentDraft.id === attachmentDraftId @@ -175,6 +203,27 @@ export const createAttachmentDraftsStoreSlice: StateCreator + _set(state => ({ + attachmentDrafts: state.attachmentDrafts.map((attachmentDraft: AttachmentDraft): AttachmentDraft => { + if (attachmentDraft.id !== attachmentDraftId) + return attachmentDraft; + + // find the removed fragments + const removedFragments = attachmentDraft.outputFragments.filter(f => !outputFragments.includes(f)); + + // remove the DBlob items associated with the removed fragments + for (let removedFragment of removedFragments) { + void removeDBlobItemFromAttachmentFragment(removedFragment); + } + + return { + ...attachmentDraft, + outputFragments, + }; + }), + })), + _getAttachment: (attachmentDraftId: AttachmentDraftId) => _get().attachmentDrafts.find(a => a.id === attachmentDraftId),