fix: image and attachments

This commit is contained in:
Louis 2024-12-13 16:35:53 +07:00
parent 5a4c5eee83
commit bb106eba01
No known key found for this signature in database
GPG Key ID: 44FA9F4D33C37DE2
25 changed files with 306 additions and 211 deletions

View File

@ -36,3 +36,10 @@ export type Assistant = {
/** Represents the metadata of the object. */
metadata?: Record<string, unknown>
}
export interface CodeInterpreterTool {
/**
* The type of tool being defined: `code_interpreter`
*/
type: 'code_interpreter'
}

View File

@ -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<Attachment> | 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<Attachment> | 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<CodeInterpreterTool | Attachment.AssistantToolsFileSearchTypeOnly>
}
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'
}

View File

@ -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}
----------------

View File

@ -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)

View File

@ -35,6 +35,7 @@ export class RetrievalTool extends InferenceTool {
await executeOnMain(
NODE,
'toolRetrievalIngestNewDocument',
data.thread?.id,
docFile,
data.model?.id,
data.model?.engine,

View File

@ -89,7 +89,9 @@ const ErrorMessage = ({ message }: { message: ThreadMessage }) => {
</span>
) : (
<>
<AutoLink text={message.content[0].text.value} />
{message?.content[0]?.text?.value && (
<AutoLink text={message?.content[0]?.text?.value} />
)}
{defaultDesc()}
</>
)}

View File

@ -8,7 +8,7 @@ import { FileInfo } from '@/types/file'
export const editPromptAtom = atom<string>('')
export const currentPromptAtom = atom<string>('')
export const fileUploadAtom = atom<FileInfo[]>([])
export const fileUploadAtom = atom<FileInfo | undefined>()
export const searchAtom = atom<string>('')

View File

@ -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([])
})
})
})

View File

@ -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({

View File

@ -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()

View File

@ -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

View File

@ -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)/)'],
})

View File

@ -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",

View File

@ -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<HTMLInputElement>) => {
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<HTMLInputElement>) => {
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 <ImageUploadPreview file={fileUpload[0].file} />
} else {
return <FileUploadPreview />
@ -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 && (
<span>
Currently, we only support 1 attachment at the same
time.

View File

@ -15,31 +15,33 @@ const FileUploadPreview = () => {
const setCurrentPrompt = useSetAtom(currentPromptAtom)
const onDeleteClick = () => {
setFileUpload([])
setFileUpload(undefined)
setCurrentPrompt('')
}
return (
<div className="flex flex-col rounded-t-lg border border-b-0 border-[hsla(var(--app-border))] p-4">
<div className="bg-secondary relative inline-flex w-60 space-x-3 rounded-lg p-4">
<Icon type={fileUpload[0].type} />
{!!fileUpload && (
<div className="bg-secondary relative inline-flex w-60 space-x-3 rounded-lg p-4">
<Icon type={fileUpload?.type} />
<div className="w-full">
<h6 className="line-clamp-1 w-3/4 truncate font-medium">
{fileUpload[0].file.name.replaceAll(/[-._]/g, ' ')}
</h6>
<p className="text-[hsla(var(--text-secondary)]">
{toGibibytes(fileUpload[0].file.size)}
</p>
</div>
<div className="w-full">
<h6 className="line-clamp-1 w-3/4 truncate font-medium">
{fileUpload?.file.name.replaceAll(/[-._]/g, ' ')}
</h6>
<p className="text-[hsla(var(--text-secondary)]">
{toGibibytes(fileUpload?.file.size)}
</p>
</div>
<div
className="absolute -right-2 -top-2 cursor-pointer rounded-full p-0.5"
onClick={onDeleteClick}
>
<XIcon size={14} className="text-background" />
<div
className="absolute -right-2 -top-2 cursor-pointer rounded-full p-0.5"
onClick={onDeleteClick}
>
<XIcon size={14} className="text-background" />
</div>
</div>
</div>
)}
</div>
)
}

View File

@ -29,7 +29,7 @@ const ImageUploadPreview: React.FC<Props> = ({ file }) => {
}
const onDeleteClick = () => {
setFileUpload([])
setFileUpload(undefined)
setCurrentPrompt('')
}

View File

@ -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 = ({
<h6 className="line-clamp-1 w-4/5 font-medium">
{name?.replaceAll(/[-._]/g, ' ')}
</h6>
<p className="text-[hsla(var(--text-secondary)]">
{/* <p className="text-[hsla(var(--text-secondary)]">
{toGibibytes(Number(size))}
</p>
</p> */}
</div>
</div>
)

View File

@ -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 (
<div className="group/image relative mb-2 inline-flex cursor-pointer overflow-hidden rounded-xl">
<div className="left-0 top-0 z-20 h-full w-full group-hover/image:inline-block">
<RelativeImage
src={annotation}
onClick={() => onViewFile(annotation)}
/>
<RelativeImage src={image} onClick={() => onViewFile(image)} />
</div>
<Tooltip
trigger={

View File

@ -33,14 +33,19 @@ const MessageContainer: React.FC<
const tokenSpeed = useAtomValue(tokenSpeedAtom)
const text = useMemo(
() => 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 (
<div className="group relative mx-auto max-w-[700px] p-4">
<div
@ -108,16 +113,8 @@ const MessageContainer: React.FC<
)}
>
<>
{messageType === ContentType.Image && (
<ImageMessage content={props.content[0]} />
)}
{messageType === ContentType.Pdf && (
<DocMessage
id={props.id}
name={props.content[0]?.text?.name}
size={props.content[0]?.text?.size}
/>
)}
{image && <ImageMessage image={image} />}
{attachedFile && <DocMessage id={props.id} name={props.id} />}
{editMessage === props.id ? (
<div>

View File

@ -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) => {

2
web/types/file.d.ts vendored
View File

@ -3,4 +3,6 @@ export type FileType = 'image' | 'pdf'
export type FileInfo = {
file: File
type: FileType
id?: string
name?: string
}

View File

@ -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
}

View File

@ -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<ChatCompletionMessage>((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,

View File

@ -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: [],
},
]);
});
},
])
})

View File

@ -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
}
}