154 lines
3.8 KiB
TypeScript
154 lines
3.8 KiB
TypeScript
"use client";
|
|
|
|
import { JSONValue } from "llamaindex";
|
|
import { useState } from "react";
|
|
import { v4 as uuidv4 } from "uuid";
|
|
import {
|
|
DocumentFile,
|
|
DocumentFileType,
|
|
MessageAnnotation,
|
|
MessageAnnotationType,
|
|
} from "..";
|
|
import { useClientConfig } from "./use-config";
|
|
|
|
const docMineTypeMap: Record<string, DocumentFileType> = {
|
|
"text/csv": "csv",
|
|
"application/pdf": "pdf",
|
|
"text/plain": "txt",
|
|
"application/vnd.openxmlformats-officedocument.wordprocessingml.document":
|
|
"docx",
|
|
};
|
|
|
|
export function useFile() {
|
|
const { backend } = useClientConfig();
|
|
const [imageUrl, setImageUrl] = useState<string | null>(null);
|
|
const [files, setFiles] = useState<DocumentFile[]>([]);
|
|
|
|
const docEqual = (a: DocumentFile, b: DocumentFile) => {
|
|
if (a.id === b.id) return true;
|
|
if (a.filename === b.filename && a.filesize === b.filesize) return true;
|
|
return false;
|
|
};
|
|
|
|
const addDoc = (file: DocumentFile) => {
|
|
const existedFile = files.find((f) => docEqual(f, file));
|
|
if (!existedFile) {
|
|
setFiles((prev) => [...prev, file]);
|
|
return true;
|
|
}
|
|
return false;
|
|
};
|
|
|
|
const removeDoc = (file: DocumentFile) => {
|
|
setFiles((prev) => prev.filter((f) => f.id !== file.id));
|
|
};
|
|
|
|
const reset = () => {
|
|
imageUrl && setImageUrl(null);
|
|
files.length && setFiles([]);
|
|
};
|
|
|
|
const uploadContent = async (
|
|
base64: string,
|
|
requestParams: any = {},
|
|
): Promise<string[]> => {
|
|
const uploadAPI = `${backend}/api/chat/upload`;
|
|
const response = await fetch(uploadAPI, {
|
|
method: "POST",
|
|
headers: {
|
|
"Content-Type": "application/json",
|
|
},
|
|
body: JSON.stringify({
|
|
base64,
|
|
...requestParams,
|
|
}),
|
|
});
|
|
if (!response.ok) throw new Error("上传文件时发生错误。");
|
|
return await response.json();
|
|
};
|
|
|
|
const getAnnotations = () => {
|
|
const annotations: MessageAnnotation[] = [];
|
|
if (imageUrl) {
|
|
annotations.push({
|
|
type: MessageAnnotationType.IMAGE,
|
|
data: { url: imageUrl },
|
|
});
|
|
}
|
|
if (files.length > 0) {
|
|
annotations.push({
|
|
type: MessageAnnotationType.DOCUMENT_FILE,
|
|
data: { files },
|
|
});
|
|
}
|
|
return annotations as JSONValue[];
|
|
};
|
|
|
|
const readContent = async (input: {
|
|
file: File;
|
|
asUrl?: boolean;
|
|
}): Promise<string> => {
|
|
const { file, asUrl } = input;
|
|
const content = await new Promise<string>((resolve, reject) => {
|
|
const reader = new FileReader();
|
|
if (asUrl) {
|
|
reader.readAsDataURL(file);
|
|
} else {
|
|
reader.readAsText(file);
|
|
}
|
|
reader.onload = () => resolve(reader.result as string);
|
|
reader.onerror = (error) => reject(error);
|
|
});
|
|
return content;
|
|
};
|
|
|
|
const uploadFile = async (file: File, requestParams: any = {}) => {
|
|
if (file.type.startsWith("image/")) {
|
|
const base64 = await readContent({ file, asUrl: true });
|
|
return setImageUrl(base64);
|
|
}
|
|
|
|
const filetype = docMineTypeMap[file.type];
|
|
if (!filetype) throw new Error("不支持的文件类型。");
|
|
const newDoc: Omit<DocumentFile, "content"> = {
|
|
id: uuidv4(),
|
|
filetype,
|
|
filename: file.name,
|
|
filesize: file.size,
|
|
};
|
|
switch (file.type) {
|
|
case "text/csv": {
|
|
const content = await readContent({ file });
|
|
return addDoc({
|
|
...newDoc,
|
|
content: {
|
|
type: "text",
|
|
value: content,
|
|
},
|
|
});
|
|
}
|
|
default: {
|
|
const base64 = await readContent({ file, asUrl: true });
|
|
const ids = await uploadContent(base64, requestParams);
|
|
return addDoc({
|
|
...newDoc,
|
|
content: {
|
|
type: "ref",
|
|
value: ids,
|
|
},
|
|
});
|
|
}
|
|
}
|
|
};
|
|
|
|
return {
|
|
imageUrl,
|
|
setImageUrl,
|
|
files,
|
|
removeDoc,
|
|
reset,
|
|
getAnnotations,
|
|
uploadFile,
|
|
};
|
|
}
|