Spaces:

Zhitomir
/

chat

Runtime error

chat / client /src /hooks /Files /useFileHandling.ts

chat init

aa2e6af over 2 years ago

10.1 kB

import { v4 } from 'uuid';
import debounce from 'lodash/debounce';
import { useQueryClient } from '@tanstack/react-query';
import { useState, useEffect, useCallback } from 'react';
import {
  megabyte,
  QueryKeys,
  EModelEndpoint,
  codeTypeMapping,
  mergeFileConfig,
  isAssistantsEndpoint,
  defaultAssistantsVersion,
  fileConfig as defaultFileConfig,
} from 'librechat-data-provider';
import type { TEndpointsConfig, TError } from 'librechat-data-provider';
import type { ExtendedFile, FileSetter } from '~/common';
import { useUploadFileMutation, useGetFileConfig } from '~/data-provider';
import { useDelayedUploadToast } from './useDelayedUploadToast';
import { useToastContext } from '~/Providers/ToastContext';
import { useChatContext } from '~/Providers/ChatContext';
import useUpdateFiles from './useUpdateFiles';

const { checkType } = defaultFileConfig;

type UseFileHandling = {
  overrideEndpoint?: EModelEndpoint;
  fileSetter?: FileSetter;
  fileFilter?: (file: File) => boolean;
  additionalMetadata?: Record<string, string | undefined>;
};

const useFileHandling = (params?: UseFileHandling) => {
  const queryClient = useQueryClient();
  const { showToast } = useToastContext();
  const [errors, setErrors] = useState<string[]>([]);
  const { startUploadTimer, clearUploadTimer } = useDelayedUploadToast();
  const { files, setFiles, setFilesLoading, conversation } = useChatContext();
  const setError = (error: string) => setErrors((prevErrors) => [...prevErrors, error]);
  const { addFile, replaceFile, updateFileById, deleteFileById } = useUpdateFiles(
    params?.fileSetter ?? setFiles,
  );

  const { data: fileConfig = defaultFileConfig } = useGetFileConfig({
    select: (data) => mergeFileConfig(data),
  });
  const endpoint =
    params?.overrideEndpoint ?? conversation?.endpointType ?? conversation?.endpoint ?? 'default';

  const { fileLimit, fileSizeLimit, totalSizeLimit, supportedMimeTypes } =
    fileConfig.endpoints[endpoint] ?? fileConfig.endpoints.default;

  const displayToast = useCallback(() => {
    if (errors.length > 1) {
      const errorList = Array.from(new Set(errors))
        .map((e, i) => `${i > 0 ? '• ' : ''}${e}\n`)
        .join('');
      showToast({
        message: errorList,
        status: 'error',
        duration: 5000,
      });
    } else if (errors.length === 1) {
      showToast({
        message: errors[0],
        status: 'error',
        duration: 5000,
      });
    }

    setErrors([]);
  }, [errors, showToast]);

  const debouncedDisplayToast = debounce(displayToast, 250);

  useEffect(() => {
    if (errors.length > 0) {
      debouncedDisplayToast();
    }

    return () => debouncedDisplayToast.cancel();
  }, [errors, debouncedDisplayToast]);

  const uploadFile = useUploadFileMutation({
    onSuccess: (data) => {
      clearUploadTimer(data.temp_file_id);
      console.log('upload success', data);
      updateFileById(
        data.temp_file_id,
        {
          progress: 0.9,
          filepath: data.filepath,
        },
        params?.additionalMetadata?.assistant_id ? true : false,
      );

      setTimeout(() => {
        updateFileById(
          data.temp_file_id,
          {
            progress: 1,
            file_id: data.file_id,
            temp_file_id: data.temp_file_id,
            filepath: data.filepath,
            type: data.type,
            height: data.height,
            width: data.width,
            filename: data.filename,
            source: data.source,
            embedded: data.embedded,
          },
          params?.additionalMetadata?.assistant_id ? true : false,
        );
      }, 300);
    },
    onError: (error, body) => {
      console.log('upload error', error);
      const file_id = body.get('file_id');
      clearUploadTimer(file_id as string);
      deleteFileById(file_id as string);
      setError(
        (error as TError)?.response?.data?.message ?? 'An error occurred while uploading the file.',
      );
    },
  });

  const startUpload = async (extendedFile: ExtendedFile) => {
    if (!endpoint) {
      setError('An error occurred while uploading the file: Endpoint is undefined');
      return;
    }

    startUploadTimer(extendedFile.file_id, extendedFile.file?.name || 'File', extendedFile.size);

    const formData = new FormData();
    formData.append(
      'file',
      extendedFile.file as File,
      encodeURIComponent(extendedFile.file?.name || 'File'),
    );
    formData.append('file_id', extendedFile.file_id);
    if (extendedFile.width) {
      formData.append('width', extendedFile.width?.toString());
    }
    if (extendedFile.height) {
      formData.append('height', extendedFile.height?.toString());
    }

    if (params?.additionalMetadata) {
      for (const [key, value] of Object.entries(params.additionalMetadata)) {
        if (value) {
          formData.append(key, value);
        }
      }
    }

    if (
      isAssistantsEndpoint(endpoint) &&
      !formData.get('assistant_id') &&
      conversation?.assistant_id
    ) {
      const endpointsConfig = queryClient.getQueryData<TEndpointsConfig>([QueryKeys.endpoints]);
      const version = endpointsConfig?.[endpoint]?.version ?? defaultAssistantsVersion[endpoint];
      formData.append('version', version);
      formData.append('assistant_id', conversation.assistant_id);
      formData.append('model', conversation?.model ?? '');
      formData.append('message_file', 'true');
    }

    formData.append('endpoint', endpoint);

    uploadFile.mutate(formData);
  };

  const validateFiles = (fileList: File[]) => {
    const existingFiles = Array.from(files.values());
    const incomingTotalSize = fileList.reduce((total, file) => total + file.size, 0);
    if (incomingTotalSize === 0) {
      setError('Empty files are not allowed.');
      return false;
    }
    const currentTotalSize = existingFiles.reduce((total, file) => total + file.size, 0);

    if (fileList.length + files.size > fileLimit) {
      setError(`You can only upload up to ${fileLimit} files at a time.`);
      return false;
    }

    for (let i = 0; i < fileList.length; i++) {
      let originalFile = fileList[i];
      let fileType = originalFile.type;
      const extension = originalFile.name.split('.').pop() ?? '';
      const knownCodeType = codeTypeMapping[extension];

      // Infer MIME type for Known Code files when the type is empty or a mismatch
      if (knownCodeType && (!fileType || fileType !== knownCodeType)) {
        fileType = knownCodeType;
      }

      // Check if the file type is still empty after the extension check
      if (!fileType) {
        setError('Unable to determine file type for: ' + originalFile.name);
        return false;
      }

      // Replace empty type with inferred type
      if (originalFile.type !== fileType) {
        const newFile = new File([originalFile], originalFile.name, { type: fileType });
        originalFile = newFile;
        fileList[i] = newFile;
      }

      if (!checkType(originalFile.type, supportedMimeTypes)) {
        console.log(originalFile);
        setError('Currently, unsupported file type: ' + originalFile.type);
        return false;
      }

      if (originalFile.size >= fileSizeLimit) {
        setError(`File size exceeds ${fileSizeLimit / megabyte} MB.`);
        return false;
      }
    }

    if (currentTotalSize + incomingTotalSize > totalSizeLimit) {
      setError(`The total size of the files cannot exceed ${totalSizeLimit / megabyte} MB.`);
      return false;
    }

    const combinedFilesInfo = [
      ...existingFiles.map(
        (file) =>
          `${file.file?.name ?? file.filename}-${file.size}-${file.type?.split('/')[0] ?? 'file'}`,
      ),
      ...fileList.map((file) => `${file.name}-${file.size}-${file.type?.split('/')[0] ?? 'file'}`),
    ];

    const uniqueFilesSet = new Set(combinedFilesInfo);

    if (uniqueFilesSet.size !== combinedFilesInfo.length) {
      setError('Duplicate file detected.');
      return false;
    }

    return true;
  };

  const loadImage = (extendedFile: ExtendedFile, preview: string) => {
    const img = new Image();
    img.onload = async () => {
      extendedFile.width = img.width;
      extendedFile.height = img.height;
      extendedFile = {
        ...extendedFile,
        progress: 0.6,
      };
      replaceFile(extendedFile);

      await startUpload(extendedFile);
      URL.revokeObjectURL(preview);
    };
    img.src = preview;
  };

  const handleFiles = async (_files: FileList | File[]) => {
    const fileList = Array.from(_files);
    /* Validate files */
    let filesAreValid: boolean;
    try {
      filesAreValid = validateFiles(fileList);
    } catch (error) {
      console.error('file validation error', error);
      setError('An error occurred while validating the file.');
      return;
    }
    if (!filesAreValid) {
      setFilesLoading(false);
      return;
    }

    /* Process files */
    for (const originalFile of fileList) {
      const file_id = v4();
      try {
        const preview = URL.createObjectURL(originalFile);
        const extendedFile: ExtendedFile = {
          file_id,
          file: originalFile,
          type: originalFile.type,
          preview,
          progress: 0.2,
          size: originalFile.size,
        };

        addFile(extendedFile);

        if (originalFile.type?.split('/')[0] === 'image') {
          loadImage(extendedFile, preview);
          continue;
        }

        await startUpload(extendedFile);
      } catch (error) {
        deleteFileById(file_id);
        console.log('file handling error', error);
        setError('An error occurred while processing the file.');
      }
    }
  };

  const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
    event.stopPropagation();
    if (event.target.files) {
      setFilesLoading(true);
      handleFiles(event.target.files);
      // reset the input
      event.target.value = '';
    }
  };

  return {
    handleFileChange,
    handleFiles,
    files,
    setFiles,
  };
};

export default useFileHandling;

	import { v4 } from 'uuid';
	import debounce from 'lodash/debounce';
	import { useQueryClient } from '@tanstack/react-query';
	import { useState, useEffect, useCallback } from 'react';
	import {
	megabyte,
	QueryKeys,
	EModelEndpoint,
	codeTypeMapping,
	mergeFileConfig,
	isAssistantsEndpoint,
	defaultAssistantsVersion,
	fileConfig as defaultFileConfig,
	} from 'librechat-data-provider';
	import type { TEndpointsConfig, TError } from 'librechat-data-provider';
	import type { ExtendedFile, FileSetter } from '~/common';
	import { useUploadFileMutation, useGetFileConfig } from '~/data-provider';
	import { useDelayedUploadToast } from './useDelayedUploadToast';
	import { useToastContext } from '~/Providers/ToastContext';
	import { useChatContext } from '~/Providers/ChatContext';
	import useUpdateFiles from './useUpdateFiles';

	const { checkType } = defaultFileConfig;

	type UseFileHandling = {
	overrideEndpoint?: EModelEndpoint;
	fileSetter?: FileSetter;
	fileFilter?: (file: File) => boolean;
	additionalMetadata?: Record<string, string \| undefined>;
	};

	const useFileHandling = (params?: UseFileHandling) => {
	const queryClient = useQueryClient();
	const { showToast } = useToastContext();
	const [errors, setErrors] = useState<string[]>([]);
	const { startUploadTimer, clearUploadTimer } = useDelayedUploadToast();
	const { files, setFiles, setFilesLoading, conversation } = useChatContext();
	const setError = (error: string) => setErrors((prevErrors) => [...prevErrors, error]);
	const { addFile, replaceFile, updateFileById, deleteFileById } = useUpdateFiles(
	params?.fileSetter ?? setFiles,
	);

	const { data: fileConfig = defaultFileConfig } = useGetFileConfig({
	select: (data) => mergeFileConfig(data),
	});
	const endpoint =
	params?.overrideEndpoint ?? conversation?.endpointType ?? conversation?.endpoint ?? 'default';

	const { fileLimit, fileSizeLimit, totalSizeLimit, supportedMimeTypes } =
	fileConfig.endpoints[endpoint] ?? fileConfig.endpoints.default;

	const displayToast = useCallback(() => {
	if (errors.length > 1) {
	const errorList = Array.from(new Set(errors))
	.map((e, i) => `${i > 0 ? '• ' : ''}${e}\n`)
	.join('');
	showToast({
	message: errorList,
	status: 'error',
	duration: 5000,
	});
	} else if (errors.length === 1) {
	showToast({
	message: errors[0],
	status: 'error',
	duration: 5000,
	});
	}

	setErrors([]);
	}, [errors, showToast]);

	const debouncedDisplayToast = debounce(displayToast, 250);

	useEffect(() => {
	if (errors.length > 0) {
	debouncedDisplayToast();
	}

	return () => debouncedDisplayToast.cancel();
	}, [errors, debouncedDisplayToast]);

	const uploadFile = useUploadFileMutation({
	onSuccess: (data) => {
	clearUploadTimer(data.temp_file_id);
	console.log('upload success', data);
	updateFileById(
	data.temp_file_id,
	{
	progress: 0.9,
	filepath: data.filepath,
	},
	params?.additionalMetadata?.assistant_id ? true : false,
	);

	setTimeout(() => {
	updateFileById(
	data.temp_file_id,
	{
	progress: 1,
	file_id: data.file_id,
	temp_file_id: data.temp_file_id,
	filepath: data.filepath,
	type: data.type,
	height: data.height,
	width: data.width,
	filename: data.filename,
	source: data.source,
	embedded: data.embedded,
	},
	params?.additionalMetadata?.assistant_id ? true : false,
	);
	}, 300);
	},
	onError: (error, body) => {
	console.log('upload error', error);
	const file_id = body.get('file_id');
	clearUploadTimer(file_id as string);
	deleteFileById(file_id as string);
	setError(
	(error as TError)?.response?.data?.message ?? 'An error occurred while uploading the file.',
	);
	},
	});

	const startUpload = async (extendedFile: ExtendedFile) => {
	if (!endpoint) {
	setError('An error occurred while uploading the file: Endpoint is undefined');
	return;
	}

	startUploadTimer(extendedFile.file_id, extendedFile.file?.name \|\| 'File', extendedFile.size);

	const formData = new FormData();
	formData.append(
	'file',
	extendedFile.file as File,
	encodeURIComponent(extendedFile.file?.name \|\| 'File'),
	);
	formData.append('file_id', extendedFile.file_id);
	if (extendedFile.width) {
	formData.append('width', extendedFile.width?.toString());
	}
	if (extendedFile.height) {
	formData.append('height', extendedFile.height?.toString());
	}

	if (params?.additionalMetadata) {
	for (const [key, value] of Object.entries(params.additionalMetadata)) {
	if (value) {
	formData.append(key, value);
	}
	}
	}

	if (
	isAssistantsEndpoint(endpoint) &&
	!formData.get('assistant_id') &&
	conversation?.assistant_id
	) {
	const endpointsConfig = queryClient.getQueryData<TEndpointsConfig>([QueryKeys.endpoints]);
	const version = endpointsConfig?.[endpoint]?.version ?? defaultAssistantsVersion[endpoint];
	formData.append('version', version);
	formData.append('assistant_id', conversation.assistant_id);
	formData.append('model', conversation?.model ?? '');
	formData.append('message_file', 'true');
	}

	formData.append('endpoint', endpoint);

	uploadFile.mutate(formData);
	};

	const validateFiles = (fileList: File[]) => {
	const existingFiles = Array.from(files.values());
	const incomingTotalSize = fileList.reduce((total, file) => total + file.size, 0);
	if (incomingTotalSize === 0) {
	setError('Empty files are not allowed.');
	return false;
	}
	const currentTotalSize = existingFiles.reduce((total, file) => total + file.size, 0);

	if (fileList.length + files.size > fileLimit) {
	setError(`You can only upload up to ${fileLimit} files at a time.`);
	return false;
	}

	for (let i = 0; i < fileList.length; i++) {
	let originalFile = fileList[i];
	let fileType = originalFile.type;
	const extension = originalFile.name.split('.').pop() ?? '';
	const knownCodeType = codeTypeMapping[extension];

	// Infer MIME type for Known Code files when the type is empty or a mismatch
	if (knownCodeType && (!fileType \|\| fileType !== knownCodeType)) {
	fileType = knownCodeType;
	}

	// Check if the file type is still empty after the extension check
	if (!fileType) {
	setError('Unable to determine file type for: ' + originalFile.name);
	return false;
	}

	// Replace empty type with inferred type
	if (originalFile.type !== fileType) {
	const newFile = new File([originalFile], originalFile.name, { type: fileType });
	originalFile = newFile;
	fileList[i] = newFile;
	}

	if (!checkType(originalFile.type, supportedMimeTypes)) {
	console.log(originalFile);
	setError('Currently, unsupported file type: ' + originalFile.type);
	return false;
	}

	if (originalFile.size >= fileSizeLimit) {
	setError(`File size exceeds ${fileSizeLimit / megabyte} MB.`);
	return false;
	}
	}

	if (currentTotalSize + incomingTotalSize > totalSizeLimit) {
	setError(`The total size of the files cannot exceed ${totalSizeLimit / megabyte} MB.`);
	return false;
	}

	const combinedFilesInfo = [
	...existingFiles.map(
	(file) =>
	`${file.file?.name ?? file.filename}-${file.size}-${file.type?.split('/')[0] ?? 'file'}`,
	),
	...fileList.map((file) => `${file.name}-${file.size}-${file.type?.split('/')[0] ?? 'file'}`),
	];

	const uniqueFilesSet = new Set(combinedFilesInfo);

	if (uniqueFilesSet.size !== combinedFilesInfo.length) {
	setError('Duplicate file detected.');
	return false;
	}

	return true;
	};

	const loadImage = (extendedFile: ExtendedFile, preview: string) => {
	const img = new Image();
	img.onload = async () => {
	extendedFile.width = img.width;
	extendedFile.height = img.height;
	extendedFile = {
	...extendedFile,
	progress: 0.6,
	};
	replaceFile(extendedFile);

	await startUpload(extendedFile);
	URL.revokeObjectURL(preview);
	};
	img.src = preview;
	};

	const handleFiles = async (_files: FileList \| File[]) => {
	const fileList = Array.from(_files);
	/* Validate files */
	let filesAreValid: boolean;
	try {
	filesAreValid = validateFiles(fileList);
	} catch (error) {
	console.error('file validation error', error);
	setError('An error occurred while validating the file.');
	return;
	}
	if (!filesAreValid) {
	setFilesLoading(false);
	return;
	}

	/* Process files */
	for (const originalFile of fileList) {
	const file_id = v4();
	try {
	const preview = URL.createObjectURL(originalFile);
	const extendedFile: ExtendedFile = {
	file_id,
	file: originalFile,
	type: originalFile.type,
	preview,
	progress: 0.2,
	size: originalFile.size,
	};

	addFile(extendedFile);

	if (originalFile.type?.split('/')[0] === 'image') {
	loadImage(extendedFile, preview);
	continue;
	}

	await startUpload(extendedFile);
	} catch (error) {
	deleteFileById(file_id);
	console.log('file handling error', error);
	setError('An error occurred while processing the file.');
	}
	}
	};

	const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
	event.stopPropagation();
	if (event.target.files) {
	setFilesLoading(true);
	handleFiles(event.target.files);
	// reset the input
	event.target.value = '';
	}
	};

	return {
	handleFileChange,
	handleFiles,
	files,
	setFiles,
	};
	};

	export default useFileHandling;