LibreChat/client/src/components/SidePanel/Builder/AssistantAvatar.tsx

import * as Popover from '@radix-ui/react-popover';
import { useState, useEffect, useRef } from 'react';
import { useQueryClient } from '@tanstack/react-query';
import {
  fileConfig as defaultFileConfig,
  QueryKeys,
  defaultOrderQuery,
  mergeFileConfig,
} from 'librechat-data-provider';
import type { UseMutationResult } from '@tanstack/react-query';
import type {
  Metadata,
  AssistantListResponse,
  Assistant,
  AssistantCreateParams,
} from 'librechat-data-provider';
import { useUploadAssistantAvatarMutation, useGetFileConfig } from '~/data-provider';
import { AssistantAvatar, NoImage, AvatarMenu } from './Images';
import { useToastContext } from '~/Providers';
// import { Spinner } from '~/components/svg';
import { useLocalize } from '~/hooks';
// import { cn } from '~/utils/';

function Avatar({
  assistant_id,
  metadata,
  createMutation,
}: {
  assistant_id: string | null;
  metadata: null | Metadata;
  createMutation: UseMutationResult<Assistant, Error, AssistantCreateParams>;
}) {
  // console.log('Avatar', assistant_id, metadata, createMutation);
  const queryClient = useQueryClient();
  const [menuOpen, setMenuOpen] = useState(false);
  const [progress, setProgress] = useState<number>(1);
  const [input, setInput] = useState<File | null>(null);
  const [previewUrl, setPreviewUrl] = useState<string | null>(null);
  const lastSeenCreatedId = useRef<string | null>(null);
  const { data: fileConfig = defaultFileConfig } = useGetFileConfig({
    select: (data) => mergeFileConfig(data),
  });

  const localize = useLocalize();
  const { showToast } = useToastContext();

  const { mutate: uploadAvatar } = useUploadAssistantAvatarMutation({
    onMutate: () => {
      setProgress(0.4);
    },
    onSuccess: (data, vars) => {
      if (!vars.postCreation) {
        showToast({ message: localize('com_ui_upload_success') });
      } else if (lastSeenCreatedId.current !== createMutation.data?.id) {
        lastSeenCreatedId.current = createMutation.data?.id ?? '';
      }

      setInput(null);
      setPreviewUrl(data.metadata?.avatar as string | null);

      const res = queryClient.getQueryData<AssistantListResponse>([
        QueryKeys.assistants,
        defaultOrderQuery,
      ]);

      if (!res?.data || !res) {
        return;
      }

      const assistants =
        res.data.map((assistant) => {
          if (assistant.id === assistant_id) {
            return {
              ...assistant,
              ...data,
            };
          }
          return assistant;
        }) ?? [];

      queryClient.setQueryData<AssistantListResponse>([QueryKeys.assistants, defaultOrderQuery], {
        ...res,
        data: assistants,
      });

      setProgress(1);
    },
    onError: (error) => {
      console.error('Error:', error);
      setInput(null);
      setPreviewUrl(null);
      showToast({ message: localize('com_ui_upload_error'), status: 'error' });
      setProgress(1);
    },
  });

  useEffect(() => {
    if (input) {
      const reader = new FileReader();
      reader.onloadend = () => {
        setPreviewUrl(reader.result as string);
      };
      reader.readAsDataURL(input);
    }
  }, [input]);

  useEffect(() => {
    setPreviewUrl((metadata?.avatar as string | undefined) ?? null);
  }, [metadata]);

  useEffect(() => {
    /** Experimental: Condition to prime avatar upload before Assistant Creation
     * - If the createMutation state Id was last seen (current) and the createMutation is successful
     * we can assume that the avatar upload has already been initiated and we can skip the upload
     *
     * The mutation state is not reset until the user deliberately selects a new assistant or an assistant is deleted
     *
     * This prevents the avatar from being uploaded multiple times before the user selects a new assistant
     * while allowing the user to upload to prime the avatar and other values before the assistant is created.
     */
    const sharedUploadCondition = !!(
      createMutation.isSuccess &&
      input &&
      previewUrl &&
      previewUrl?.includes('base64')
    );
    if (sharedUploadCondition && lastSeenCreatedId.current === createMutation.data?.id) {
      return;
    }

    if (sharedUploadCondition && createMutation.data.id) {
      console.log('[AssistantAvatar] Uploading Avatar after Assistant Creation');

      const formData = new FormData();
      formData.append('file', input, input.name);
      formData.append('assistant_id', createMutation.data.id);

      if (typeof createMutation.data?.metadata === 'object') {
        formData.append('metadata', JSON.stringify(createMutation.data?.metadata));
      }

      uploadAvatar({
        assistant_id: createMutation.data.id,
        postCreation: true,
        formData,
      });
    }
  }, [createMutation.data, createMutation.isSuccess, input, previewUrl, uploadAvatar]);

  const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>): void => {
    const file = event.target.files?.[0];

    if (fileConfig.avatarSizeLimit && file && file.size <= fileConfig.avatarSizeLimit) {
      if (!file) {
        console.error('No file selected');
        return;
      }

      setInput(file);
      setMenuOpen(false);

      if (!assistant_id) {
        // wait for successful form submission before uploading avatar
        console.log('[AssistantAvatar] No assistant_id, will wait until form submission + upload');
        return;
      }

      const formData = new FormData();
      formData.append('file', file, file.name);
      formData.append('assistant_id', assistant_id);

      if (typeof metadata === 'object') {
        formData.append('metadata', JSON.stringify(metadata));
      }

      uploadAvatar({
        assistant_id,
        formData,
      });
    } else {
      showToast({
        message: localize('com_ui_upload_invalid'),
        status: 'error',
      });
    }

    setMenuOpen(false);
  };

  return (
    <Popover.Root open={menuOpen} onOpenChange={setMenuOpen}>
      <div className="flex w-full items-center justify-center gap-4">
        <Popover.Trigger asChild>
          <button type="button" className="h-20 w-20">
            {previewUrl ? <AssistantAvatar url={previewUrl} progress={progress} /> : <NoImage />}
          </button>
        </Popover.Trigger>
      </div>
      {<AvatarMenu handleFileChange={handleFileChange} />}
    </Popover.Root>
  );
}

export default Avatar;