From bb106eba011bc2bd95691afa6afee0f2e0707669 Mon Sep 17 00:00:00 2001 From: Louis Date: Fri, 13 Dec 2024 16:35:53 +0700 Subject: [PATCH] fix: image and attachments --- core/src/types/assistant/assistantEntity.ts | 7 + core/src/types/message/messageEntity.ts | 56 +++++- extensions/assistant-extension/src/index.ts | 2 +- .../assistant-extension/src/node/index.ts | 5 +- .../src/tools/retrieval.ts | 1 + web/containers/ErrorMessage/index.tsx | 4 +- web/containers/Providers/Jotai.tsx | 2 +- web/helpers/atoms/Model.atom.test.ts | 12 +- web/hooks/useCreateNewThread.ts | 3 +- web/hooks/useDropModelBinaries.test.ts | 4 + web/hooks/useSendChatMessage.ts | 14 +- web/jest.config.js | 2 +- web/package.json | 3 + .../ThreadCenterPanel/ChatInput/index.tsx | 26 ++- .../FileUploadPreview/index.tsx | 36 ++-- .../ImageUploadPreview/index.tsx | 2 +- .../TextMessage/DocMessage.tsx | 14 +- .../TextMessage/ImageMessage.tsx | 15 +- .../ThreadCenterPanel/TextMessage/index.tsx | 23 +-- .../Thread/ThreadCenterPanel/index.tsx | 20 +- web/types/file.d.ts | 2 + web/utils/file.ts | 16 ++ web/utils/messageRequestBuilder.ts | 17 +- web/utils/threadMessageBuilder.test.ts | 182 +++++++++--------- web/utils/threadMessageBuilder.ts | 49 ++--- 25 files changed, 306 insertions(+), 211 deletions(-) diff --git a/core/src/types/assistant/assistantEntity.ts b/core/src/types/assistant/assistantEntity.ts index 27592e26b..42617a4b5 100644 --- a/core/src/types/assistant/assistantEntity.ts +++ b/core/src/types/assistant/assistantEntity.ts @@ -36,3 +36,10 @@ export type Assistant = { /** Represents the metadata of the object. */ metadata?: Record } + +export interface CodeInterpreterTool { + /** + * The type of tool being defined: `code_interpreter` + */ + type: 'code_interpreter' +} diff --git a/core/src/types/message/messageEntity.ts b/core/src/types/message/messageEntity.ts index b74176fd4..7c2774da6 100644 --- a/core/src/types/message/messageEntity.ts +++ b/core/src/types/message/messageEntity.ts @@ -1,3 +1,4 @@ +import { CodeInterpreterTool } from '../assistant' import { ChatCompletionMessage, ChatCompletionRole } from '../inference' import { ModelInfo } from '../model' import { Thread } from '../thread' @@ -15,6 +16,10 @@ export type ThreadMessage = { thread_id: string /** The assistant id of this thread. **/ assistant_id?: string + /** + * A list of files attached to the message, and the tools they were added to. + */ + attachments?: Array | null /** The role of the author of this message. **/ role: ChatCompletionRole /** The content of this message. **/ @@ -52,6 +57,11 @@ export type MessageRequest = { */ assistantId?: string + /** + * A list of files attached to the message, and the tools they were added to. + */ + attachments: Array | null + /** Messages for constructing a chat completion request **/ messages?: ChatCompletionMessage[] @@ -98,7 +108,6 @@ export enum ErrorCode { export enum ContentType { Text = 'text', Image = 'image_url', - Pdf = 'pdf', } /** @@ -108,8 +117,15 @@ export enum ContentType { export type ContentValue = { value: string annotations: string[] - name?: string - size?: number +} + +/** + * The `ImageContentValue` type defines the shape of a content value object of image type + * @data_transfer_object + */ +export type ImageContentValue = { + detail?: string + url?: string } /** @@ -118,5 +134,37 @@ export type ContentValue = { */ export type ThreadContent = { type: ContentType - text: ContentValue + text?: ContentValue + image_url?: ImageContentValue +} + +export interface Attachment { + /** + * The ID of the file to attach to the message. + */ + file_id?: string + + /** + * The tools to add this file to. + */ + tools?: Array +} + +export namespace Attachment { + export interface AssistantToolsFileSearchTypeOnly { + /** + * The type of tool being defined: `file_search` + */ + type: 'file_search' + } +} + +/** + * On an incomplete message, details about why the message is incomplete. + */ +export interface IncompleteDetails { + /** + * The reason the message is incomplete. + */ + reason: 'content_filter' | 'max_tokens' | 'run_cancelled' | 'run_expired' | 'run_failed' } diff --git a/extensions/assistant-extension/src/index.ts b/extensions/assistant-extension/src/index.ts index 6705483d6..0b3a1ec40 100644 --- a/extensions/assistant-extension/src/index.ts +++ b/extensions/assistant-extension/src/index.ts @@ -141,7 +141,7 @@ export default class JanAssistantExtension extends AssistantExtension { top_k: 2, chunk_size: 1024, chunk_overlap: 64, - retrieval_template: `Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. + retrieval_template: `Use the following pieces of context to answer the question at the end. ---------------- CONTEXT: {CONTEXT} ---------------- diff --git a/extensions/assistant-extension/src/node/index.ts b/extensions/assistant-extension/src/node/index.ts index 83a4a1983..11e8f49c4 100644 --- a/extensions/assistant-extension/src/node/index.ts +++ b/extensions/assistant-extension/src/node/index.ts @@ -9,13 +9,14 @@ export function toolRetrievalUpdateTextSplitter( retrieval.updateTextSplitter(chunkSize, chunkOverlap) } export async function toolRetrievalIngestNewDocument( + thread: string, file: string, model: string, engine: string, useTimeWeighted: boolean ) { - const filePath = path.join(getJanDataFolderPath(), normalizeFilePath(file)) - const threadPath = path.dirname(filePath.replace('files', '')) + const threadPath = path.join(getJanDataFolderPath(), 'threads', thread) + const filePath = path.join(getJanDataFolderPath(), 'files', file) retrieval.updateEmbeddingEngine(model, engine) return retrieval .ingestAgentKnowledge(filePath, `${threadPath}/memory`, useTimeWeighted) diff --git a/extensions/assistant-extension/src/tools/retrieval.ts b/extensions/assistant-extension/src/tools/retrieval.ts index 763192287..b1a0c3cba 100644 --- a/extensions/assistant-extension/src/tools/retrieval.ts +++ b/extensions/assistant-extension/src/tools/retrieval.ts @@ -35,6 +35,7 @@ export class RetrievalTool extends InferenceTool { await executeOnMain( NODE, 'toolRetrievalIngestNewDocument', + data.thread?.id, docFile, data.model?.id, data.model?.engine, diff --git a/web/containers/ErrorMessage/index.tsx b/web/containers/ErrorMessage/index.tsx index 2af6065a6..95b87fc53 100644 --- a/web/containers/ErrorMessage/index.tsx +++ b/web/containers/ErrorMessage/index.tsx @@ -89,7 +89,9 @@ const ErrorMessage = ({ message }: { message: ThreadMessage }) => { ) : ( <> - + {message?.content[0]?.text?.value && ( + + )} {defaultDesc()} )} diff --git a/web/containers/Providers/Jotai.tsx b/web/containers/Providers/Jotai.tsx index 8f1433ea0..5371097f4 100644 --- a/web/containers/Providers/Jotai.tsx +++ b/web/containers/Providers/Jotai.tsx @@ -8,7 +8,7 @@ import { FileInfo } from '@/types/file' export const editPromptAtom = atom('') export const currentPromptAtom = atom('') -export const fileUploadAtom = atom([]) +export const fileUploadAtom = atom() export const searchAtom = atom('') diff --git a/web/helpers/atoms/Model.atom.test.ts b/web/helpers/atoms/Model.atom.test.ts index 923f24df4..b4eb87e7a 100644 --- a/web/helpers/atoms/Model.atom.test.ts +++ b/web/helpers/atoms/Model.atom.test.ts @@ -58,7 +58,9 @@ describe('Model.atom.ts', () => { setAtom.current({ id: '1' } as any) }) expect(getAtom.current).toEqual([{ id: '1' }]) - reset.current([]) + act(() => { + reset.current([]) + }) }) }) @@ -83,7 +85,9 @@ describe('Model.atom.ts', () => { removeAtom.current('1') }) expect(getAtom.current).toEqual([]) - reset.current([]) + act(() => { + reset.current([]) + }) }) }) @@ -113,7 +117,9 @@ describe('Model.atom.ts', () => { removeAtom.current('1') }) expect(getAtom.current).toEqual([]) - reset.current([]) + act(() => { + reset.current([]) + }) }) }) diff --git a/web/hooks/useCreateNewThread.ts b/web/hooks/useCreateNewThread.ts index ec029c2df..5ecfb649a 100644 --- a/web/hooks/useCreateNewThread.ts +++ b/web/hooks/useCreateNewThread.ts @@ -20,6 +20,7 @@ import { fileUploadAtom } from '@/containers/Providers/Jotai' import { toaster } from '@/containers/Toast' import { isLocalEngine } from '@/utils/modelEngine' + import { useActiveModel } from './useActiveModel' import useRecommendedModel from './useRecommendedModel' @@ -168,7 +169,7 @@ export const useCreateNewThread = () => { }) // Delete the file upload state - setFileUpload([]) + setFileUpload(undefined) setActiveThread(createdThread) } catch (ex) { return toaster({ diff --git a/web/hooks/useDropModelBinaries.test.ts b/web/hooks/useDropModelBinaries.test.ts index dad8c6178..7ca5a479e 100644 --- a/web/hooks/useDropModelBinaries.test.ts +++ b/web/hooks/useDropModelBinaries.test.ts @@ -1,3 +1,6 @@ +/** + * @jest-environment jsdom + */ // useDropModelBinaries.test.ts import { renderHook, act } from '@testing-library/react' @@ -18,6 +21,7 @@ jest.mock('jotai', () => ({ jest.mock('uuid') jest.mock('@/utils/file') jest.mock('@/containers/Toast') +jest.mock("@uppy/core") describe('useDropModelBinaries', () => { const mockSetImportingModels = jest.fn() diff --git a/web/hooks/useSendChatMessage.ts b/web/hooks/useSendChatMessage.ts index 86c2e41b6..bbe5e3cd7 100644 --- a/web/hooks/useSendChatMessage.ts +++ b/web/hooks/useSendChatMessage.ts @@ -134,11 +134,9 @@ export default function useSendChatMessage() { setCurrentPrompt('') setEditPrompt('') - let base64Blob = fileUpload[0] - ? await getBase64(fileUpload[0].file) - : undefined + let base64Blob = fileUpload ? await getBase64(fileUpload.file) : undefined - if (base64Blob && fileUpload[0]?.type === 'image') { + if (base64Blob && fileUpload?.type === 'image') { // Compress image base64Blob = await compressImage(base64Blob, 512) } @@ -171,7 +169,7 @@ export default function useSendChatMessage() { ).addSystemMessage(activeAssistantRef.current?.instructions) if (!isResend) { - requestBuilder.pushMessage(prompt, base64Blob, fileUpload[0]?.type) + requestBuilder.pushMessage(prompt, base64Blob, fileUpload) // Build Thread Message to persist const threadMessageBuilder = new ThreadMessageBuilder( @@ -207,7 +205,7 @@ export default function useSendChatMessage() { selectedModelRef.current?.id ?? activeAssistantRef.current?.model.id if (base64Blob) { - setFileUpload([]) + setFileUpload(undefined) } if (modelRef.current?.id !== modelId && modelId) { @@ -222,9 +220,7 @@ export default function useSendChatMessage() { // Process message request with Assistants tools const request = await ToolManager.instance().process( requestBuilder.build(), - activeThreadRef.current.assistants?.flatMap( - (assistant) => assistant.tools ?? [] - ) ?? [] + activeAssistantRef?.current.tools ?? [] ) // Request for inference diff --git a/web/jest.config.js b/web/jest.config.js index f78007532..27e8d0bda 100644 --- a/web/jest.config.js +++ b/web/jest.config.js @@ -37,5 +37,5 @@ const config = { // module.exports = createJestConfig(config) module.exports = async () => ({ ...(await createJestConfig(config)()), - transformIgnorePatterns: ['/node_modules/(?!(layerr)/)'], + transformIgnorePatterns: ['/node_modules/(?!(layerr|nanoid|@uppy|preact)/)'], }) diff --git a/web/package.json b/web/package.json index 13f646a6f..6aa06dc8e 100644 --- a/web/package.json +++ b/web/package.json @@ -17,6 +17,9 @@ "@janhq/core": "link:./core", "@janhq/joi": "link:./joi", "@tanstack/react-virtual": "^3.10.9", + "@uppy/core": "^4.3.0", + "@uppy/react": "^4.0.4", + "@uppy/xhr-upload": "^4.2.3", "autoprefixer": "10.4.16", "class-variance-authority": "^0.7.0", "framer-motion": "^10.16.4", diff --git a/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx index b88b26732..b3246a26b 100644 --- a/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx +++ b/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx @@ -24,6 +24,7 @@ import { useActiveModel } from '@/hooks/useActiveModel' import useSendChatMessage from '@/hooks/useSendChatMessage' +import { uploader } from '@/utils/file' import { isLocalEngine } from '@/utils/modelEngine' import FileUploadPreview from '../FileUploadPreview' @@ -71,6 +72,7 @@ const ChatInput = () => { const activeAssistant = useAtomValue(activeAssistantAtom) const { stopInference } = useActiveModel() + const upload = uploader() const [activeTabThreadRightPanel, setActiveTabThreadRightPanel] = useAtom( activeTabThreadRightPanelAtom ) @@ -104,18 +106,26 @@ const ChatInput = () => { const handleFileChange = (event: React.ChangeEvent) => { const file = event.target.files?.[0] if (!file) return - setFileUpload([{ file: file, type: 'pdf' }]) + upload.addFile(file) + upload.upload().then((data) => { + setFileUpload({ + file: file, + type: 'pdf', + id: data?.successful?.[0]?.response?.body?.id, + name: data?.successful?.[0]?.response?.body?.filename, + }) + }) } const handleImageChange = (event: React.ChangeEvent) => { const file = event.target.files?.[0] if (!file) return - setFileUpload([{ file: file, type: 'image' }]) + setFileUpload({ file: file, type: 'image' }) } const renderPreview = (fileUpload: any) => { - if (fileUpload.length > 0) { - if (fileUpload[0].type === 'image') { + if (fileUpload) { + if (fileUpload.type === 'image') { return } else { return @@ -132,7 +142,7 @@ const ChatInput = () => { 'relative mb-1 max-h-[400px] resize-none rounded-lg border border-[hsla(var(--app-border))] p-3 pr-20', 'focus-within:outline-none focus-visible:outline-0 focus-visible:ring-1 focus-visible:ring-[hsla(var(--primary-bg))] focus-visible:ring-offset-0', 'overflow-y-auto', - fileUpload.length && 'rounded-t-none', + fileUpload && 'rounded-t-none', experimentalFeature && 'pl-10', activeSettingInputBox && 'pb-14 pr-16' )} @@ -154,7 +164,7 @@ const ChatInput = () => { className="absolute left-3 top-2.5" onClick={(e) => { if ( - fileUpload.length > 0 || + !!fileUpload || (activeAssistant?.tools && !activeAssistant?.tools[0]?.enabled && !activeAssistant?.model.settings?.vision_model) @@ -178,12 +188,12 @@ const ChatInput = () => { } content={ <> - {fileUpload.length > 0 || + {!!fileUpload || (activeAssistant?.tools && !activeAssistant?.tools[0]?.enabled && !activeAssistant?.model.settings?.vision_model && ( <> - {fileUpload.length !== 0 && ( + {!!fileUpload && ( Currently, we only support 1 attachment at the same time. diff --git a/web/screens/Thread/ThreadCenterPanel/FileUploadPreview/index.tsx b/web/screens/Thread/ThreadCenterPanel/FileUploadPreview/index.tsx index 348e915e6..0e4872e10 100644 --- a/web/screens/Thread/ThreadCenterPanel/FileUploadPreview/index.tsx +++ b/web/screens/Thread/ThreadCenterPanel/FileUploadPreview/index.tsx @@ -15,31 +15,33 @@ const FileUploadPreview = () => { const setCurrentPrompt = useSetAtom(currentPromptAtom) const onDeleteClick = () => { - setFileUpload([]) + setFileUpload(undefined) setCurrentPrompt('') } return (
-
- + {!!fileUpload && ( +
+ -
-
- {fileUpload[0].file.name.replaceAll(/[-._]/g, ' ')} -
-

- {toGibibytes(fileUpload[0].file.size)} -

-
+
+
+ {fileUpload?.file.name.replaceAll(/[-._]/g, ' ')} +
+

+ {toGibibytes(fileUpload?.file.size)} +

+
-
- +
+ +
-
+ )}
) } diff --git a/web/screens/Thread/ThreadCenterPanel/ImageUploadPreview/index.tsx b/web/screens/Thread/ThreadCenterPanel/ImageUploadPreview/index.tsx index b43b80830..7fa9e417a 100644 --- a/web/screens/Thread/ThreadCenterPanel/ImageUploadPreview/index.tsx +++ b/web/screens/Thread/ThreadCenterPanel/ImageUploadPreview/index.tsx @@ -29,7 +29,7 @@ const ImageUploadPreview: React.FC = ({ file }) => { } const onDeleteClick = () => { - setFileUpload([]) + setFileUpload(undefined) setCurrentPrompt('') } diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/DocMessage.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/DocMessage.tsx index 9c0289734..69d61d0d5 100644 --- a/web/screens/Thread/ThreadCenterPanel/TextMessage/DocMessage.tsx +++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/DocMessage.tsx @@ -11,15 +11,7 @@ import { openFileTitle } from '@/utils/titleUtils' import Icon from '../FileUploadPreview/Icon' -const DocMessage = ({ - id, - name, - size, -}: { - id: string - name?: string - size?: number -}) => { +const DocMessage = ({ id, name }: { id: string; name?: string }) => { const { onViewFile, onViewFileContainer } = usePath() return ( @@ -44,9 +36,9 @@ const DocMessage = ({
{name?.replaceAll(/[-._]/g, ' ')}
-

+ {/*

{toGibibytes(Number(size))} -

+

*/}
) diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/ImageMessage.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/ImageMessage.tsx index 117f259c0..e83d35fbb 100644 --- a/web/screens/Thread/ThreadCenterPanel/TextMessage/ImageMessage.tsx +++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/ImageMessage.tsx @@ -1,6 +1,5 @@ -import { memo, useMemo } from 'react' +import { memo } from 'react' -import { ThreadContent } from '@janhq/core' import { Tooltip } from '@janhq/joi' import { FolderOpenIcon } from 'lucide-react' @@ -11,21 +10,13 @@ import { openFileTitle } from '@/utils/titleUtils' import { RelativeImage } from '../TextMessage/RelativeImage' -const ImageMessage = ({ content }: { content: ThreadContent }) => { +const ImageMessage = ({ image }: { image: string }) => { const { onViewFile, onViewFileContainer } = usePath() - const annotation = useMemo( - () => content?.text?.annotations[0] ?? '', - [content] - ) - return (
- onViewFile(annotation)} - /> + onViewFile(image)} />
props.content[0]?.text?.value ?? '', + () => + props.content.find((e) => e.type === ContentType.Text)?.text?.value ?? '', [props.content] ) - const messageType = useMemo( - () => props.content[0]?.type ?? '', + + const image = useMemo( + () => + props.content.find((e) => e.type === ContentType.Image)?.image_url?.url, [props.content] ) + const attachedFile = useMemo(() => 'attachments' in props, [props]) + return (
<> - {messageType === ContentType.Image && ( - - )} - {messageType === ContentType.Pdf && ( - - )} + {image && } + {attachedFile && } {editMessage === props.id ? (
diff --git a/web/screens/Thread/ThreadCenterPanel/index.tsx b/web/screens/Thread/ThreadCenterPanel/index.tsx index c9a92f455..ca04f9e59 100644 --- a/web/screens/Thread/ThreadCenterPanel/index.tsx +++ b/web/screens/Thread/ThreadCenterPanel/index.tsx @@ -22,6 +22,8 @@ import { reloadModelAtom } from '@/hooks/useSendChatMessage' import ChatBody from '@/screens/Thread/ThreadCenterPanel/ChatBody' +import { uploader } from '@/utils/file' + import ChatInput from './ChatInput' import RequestDownloadModel from './RequestDownloadModel' @@ -57,7 +59,7 @@ const ThreadCenterPanel = () => { const experimentalFeature = useAtomValue(experimentalFeatureEnabledAtom) const activeThread = useAtomValue(activeThreadAtom) const activeAssistant = useAtomValue(activeAssistantAtom) - + const upload = uploader() const acceptedFormat: Accept = activeAssistant?.model.settings?.vision_model ? { 'application/pdf': ['.pdf'], @@ -93,7 +95,7 @@ const ThreadCenterPanel = () => { } }, onDragLeave: () => setDragOver(false), - onDrop: (files, rejectFiles) => { + onDrop: async (files, rejectFiles) => { // Retrieval file drag and drop is experimental feature if (!experimentalFeature) return if ( @@ -106,7 +108,19 @@ const ThreadCenterPanel = () => { ) return const imageType = files[0]?.type.includes('image') - setFileUpload([{ file: files[0], type: imageType ? 'image' : 'pdf' }]) + if (imageType) { + setFileUpload({ file: files[0], type: 'image' }) + } else { + upload.addFile(files[0]) + upload.upload().then((data) => { + setFileUpload({ + file: files[0], + type: imageType ? 'image' : 'pdf', + id: data?.successful?.[0]?.response?.body?.id, + name: data?.successful?.[0]?.response?.body?.filename, + }) + }) + } setDragOver(false) }, onDropRejected: (e) => { diff --git a/web/types/file.d.ts b/web/types/file.d.ts index 737c5e380..aac6ba379 100644 --- a/web/types/file.d.ts +++ b/web/types/file.d.ts @@ -3,4 +3,6 @@ export type FileType = 'image' | 'pdf' export type FileInfo = { file: File type: FileType + id?: string + name?: string } diff --git a/web/utils/file.ts b/web/utils/file.ts index 4a14de247..2eb6d9f64 100644 --- a/web/utils/file.ts +++ b/web/utils/file.ts @@ -1,4 +1,6 @@ import { baseName } from '@janhq/core' +import Uppy from '@uppy/core' +import XHR from '@uppy/xhr-upload' export type FilePathWithSize = { path: string @@ -27,3 +29,17 @@ export const getFileInfoFromFile = async ( } return result } + +export const uploader = () => { + const uppy = new Uppy().use(XHR, { + endpoint: 'http://127.0.0.1:39291/v1/files', + method: 'POST', + fieldName: 'file', + formData: true, + limit: 1, + }) + uppy.setMeta({ + purpose: 'assistants', + }) + return uppy +} diff --git a/web/utils/messageRequestBuilder.ts b/web/utils/messageRequestBuilder.ts index 63b14d769..99560636e 100644 --- a/web/utils/messageRequestBuilder.ts +++ b/web/utils/messageRequestBuilder.ts @@ -15,7 +15,7 @@ import { ulid } from 'ulidx' import { Stack } from '@/utils/Stack' -import { FileType } from '@/types/file' +import { FileInfo, FileType } from '@/types/file' export class MessageRequestBuilder { msgId: string @@ -38,7 +38,7 @@ export class MessageRequestBuilder { .filter((e) => e.status !== MessageStatus.Error) .map((msg) => ({ role: msg.role, - content: msg.content[0]?.text.value ?? '.', + content: msg.content[0]?.text?.value ?? '.', })) } @@ -46,11 +46,11 @@ export class MessageRequestBuilder { pushMessage( message: string, base64Blob: string | undefined, - fileContentType: FileType + fileInfo?: FileInfo ) { - if (base64Blob && fileContentType === 'pdf') - return this.addDocMessage(message) - else if (base64Blob && fileContentType === 'image') { + if (base64Blob && fileInfo?.type === 'pdf') + return this.addDocMessage(message, fileInfo?.name) + else if (base64Blob && fileInfo?.type === 'image') { return this.addImageMessage(message, base64Blob) } this.messages = [ @@ -77,7 +77,7 @@ export class MessageRequestBuilder { } // Chainable - addDocMessage(prompt: string) { + addDocMessage(prompt: string, name?: string) { const message: ChatCompletionMessage = { role: ChatCompletionRole.User, content: [ @@ -88,7 +88,7 @@ export class MessageRequestBuilder { { type: ChatCompletionMessageContentType.Doc, doc_url: { - url: `threads/${this.thread.id}/files/${this.msgId}.pdf`, + url: name ?? `${this.msgId}.pdf`, }, }, ] as ChatCompletionMessageContent, @@ -163,6 +163,7 @@ export class MessageRequestBuilder { return { id: this.msgId, type: this.type, + attachments: [], threadId: this.thread.id, messages: this.normalizeMessages(this.messages), model: this.model, diff --git a/web/utils/threadMessageBuilder.test.ts b/web/utils/threadMessageBuilder.test.ts index cc192a5c1..856c9ad06 100644 --- a/web/utils/threadMessageBuilder.test.ts +++ b/web/utils/threadMessageBuilder.test.ts @@ -1,100 +1,100 @@ +import { + ChatCompletionRole, + MessageRequestType, + MessageStatus, +} from '@janhq/core' -import { ChatCompletionRole, MessageStatus } from '@janhq/core' +import { ThreadMessageBuilder } from './threadMessageBuilder' +import { MessageRequestBuilder } from './messageRequestBuilder' - import { ThreadMessageBuilder } from './threadMessageBuilder' - import { MessageRequestBuilder } from './messageRequestBuilder' - -import { ContentType } from '@janhq/core'; - describe('ThreadMessageBuilder', () => { - it('testBuildMethod', () => { - const msgRequest = new MessageRequestBuilder( - 'type', - { model: 'model' }, - { id: 'thread-id' }, - [] - ) - const builder = new ThreadMessageBuilder(msgRequest) - const result = builder.build() - - expect(result.id).toBe(msgRequest.msgId) - expect(result.thread_id).toBe(msgRequest.thread.id) - expect(result.role).toBe(ChatCompletionRole.User) - expect(result.status).toBe(MessageStatus.Ready) - expect(result.created).toBeDefined() - expect(result.updated).toBeDefined() - expect(result.object).toBe('thread.message') - expect(result.content).toEqual([]) - }) +import { ContentType } from '@janhq/core' +describe('ThreadMessageBuilder', () => { + it('testBuildMethod', () => { + const msgRequest = new MessageRequestBuilder( + MessageRequestType.Thread, + { model: 'model' } as any, + { id: 'thread-id' } as any, + [] + ) + const builder = new ThreadMessageBuilder(msgRequest) + const result = builder.build() + + expect(result.id).toBe(msgRequest.msgId) + expect(result.thread_id).toBe(msgRequest.thread.id) + expect(result.role).toBe(ChatCompletionRole.User) + expect(result.status).toBe(MessageStatus.Ready) + expect(result.created).toBeDefined() + expect(result.updated).toBeDefined() + expect(result.object).toBe('thread.message') + expect(result.content).toEqual([]) }) +}) - it('testPushMessageWithPromptOnly', () => { - const msgRequest = new MessageRequestBuilder( - 'type', - { model: 'model' }, - { id: 'thread-id' }, - [] - ); - const builder = new ThreadMessageBuilder(msgRequest); - const prompt = 'test prompt'; - builder.pushMessage(prompt, undefined, []); - expect(builder.content).toEqual([ - { - type: ContentType.Text, - text: { - value: prompt, - annotations: [], - }, +it('testPushMessageWithPromptOnly', () => { + const msgRequest = new MessageRequestBuilder( + MessageRequestType.Thread, + { model: 'model' } as any, + { id: 'thread-id' } as any, + [] + ) + const builder = new ThreadMessageBuilder(msgRequest) + const prompt = 'test prompt' + builder.pushMessage(prompt, undefined, undefined) + expect(builder.content).toEqual([ + { + type: ContentType.Text, + text: { + value: prompt, + annotations: [], }, - ]); - }); + }, + ]) +}) - - it('testPushMessageWithPdf', () => { - const msgRequest = new MessageRequestBuilder( - 'type', - { model: 'model' }, - { id: 'thread-id' }, - [] - ); - const builder = new ThreadMessageBuilder(msgRequest); - const prompt = 'test prompt'; - const base64 = 'test base64'; - const fileUpload = [{ type: 'pdf', file: { name: 'test.pdf', size: 1000 } }]; - builder.pushMessage(prompt, base64, fileUpload); - expect(builder.content).toEqual([ - { - type: ContentType.Pdf, - text: { - value: prompt, - annotations: [base64], - name: fileUpload[0].file.name, - size: fileUpload[0].file.size, - }, +it('testPushMessageWithPdf', () => { + const msgRequest = new MessageRequestBuilder( + MessageRequestType.Thread, + { model: 'model' } as any, + { id: 'thread-id' } as any, + [] + ) + const builder = new ThreadMessageBuilder(msgRequest) + const prompt = 'test prompt' + const base64 = 'test base64' + const fileUpload = [ + { type: 'pdf', file: { name: 'test.pdf', size: 1000 } }, + ] as any + builder.pushMessage(prompt, base64, fileUpload) + expect(builder.content).toEqual([ + { + type: ContentType.Text, + text: { + value: prompt, + annotations: [], }, - ]); - }); + }, + ]) +}) - - it('testPushMessageWithImage', () => { - const msgRequest = new MessageRequestBuilder( - 'type', - { model: 'model' }, - { id: 'thread-id' }, - [] - ); - const builder = new ThreadMessageBuilder(msgRequest); - const prompt = 'test prompt'; - const base64 = 'test base64'; - const fileUpload = [{ type: 'image', file: { name: 'test.jpg', size: 1000 } }]; - builder.pushMessage(prompt, base64, fileUpload); - expect(builder.content).toEqual([ - { - type: ContentType.Image, - text: { - value: prompt, - annotations: [base64], - }, +it('testPushMessageWithImage', () => { + const msgRequest = new MessageRequestBuilder( + MessageRequestType.Thread, + { model: 'model' } as any, + { id: 'thread-id' } as any, + [] + ) + const builder = new ThreadMessageBuilder(msgRequest) + const prompt = 'test prompt' + const base64 = 'test base64' + const fileUpload = [{ type: 'image', file: { name: 'test.jpg', size: 1000 } }] + builder.pushMessage(prompt, base64, fileUpload as any) + expect(builder.content).toEqual([ + { + type: ContentType.Text, + text: { + value: prompt, + annotations: [], }, - ]); - }); - + }, + ]) +}) diff --git a/web/utils/threadMessageBuilder.ts b/web/utils/threadMessageBuilder.ts index 1f55e4d2d..d1aac2f84 100644 --- a/web/utils/threadMessageBuilder.ts +++ b/web/utils/threadMessageBuilder.ts @@ -1,4 +1,5 @@ import { + Attachment, ChatCompletionRole, ContentType, MessageStatus, @@ -14,6 +15,7 @@ export class ThreadMessageBuilder { messageRequest: MessageRequestBuilder content: ThreadContent[] = [] + attachments: Attachment[] = [] constructor(messageRequest: MessageRequestBuilder) { this.messageRequest = messageRequest @@ -24,6 +26,7 @@ export class ThreadMessageBuilder { return { id: this.messageRequest.msgId, thread_id: this.messageRequest.thread.id, + attachments: this.attachments, role: ChatCompletionRole.User, status: MessageStatus.Ready, created: timestamp, @@ -36,31 +39,9 @@ export class ThreadMessageBuilder { pushMessage( prompt: string, base64: string | undefined, - fileUpload: FileInfo[] + fileUpload?: FileInfo ) { - if (base64 && fileUpload[0]?.type === 'image') { - this.content.push({ - type: ContentType.Image, - text: { - value: prompt, - annotations: [base64], - }, - }) - } - - if (base64 && fileUpload[0]?.type === 'pdf') { - this.content.push({ - type: ContentType.Pdf, - text: { - value: prompt, - annotations: [base64], - name: fileUpload[0].file.name, - size: fileUpload[0].file.size, - }, - }) - } - - if (prompt && !base64) { + if (prompt) { this.content.push({ type: ContentType.Text, text: { @@ -69,6 +50,26 @@ export class ThreadMessageBuilder { }, }) } + if (base64 && fileUpload?.type === 'image') { + this.content.push({ + type: ContentType.Image, + image_url: { + url: base64, + }, + }) + } + + if (base64 && fileUpload?.type === 'pdf') { + this.attachments.push({ + file_id: fileUpload.id, + tools: [ + { + type: 'file_search', + }, + ], + }) + } + return this } }