Skip to content

Commit

Permalink
feat(ts): add formatChat util
Browse files Browse the repository at this point in the history
  • Loading branch information
jhen0409 committed Jul 28, 2024
1 parent c495df8 commit a75662c
Show file tree
Hide file tree
Showing 2 changed files with 108 additions and 0 deletions.
61 changes: 61 additions & 0 deletions src/__tests__/chat.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
import { formatChat } from '../chat'

describe('formatChat', () => {
it('should format chat messages', () => {
const messages = [
{
role: 'user',
content: 'Hello, world!',
},
{
role: 'bot',
content: [
{
text: 'Hello, user!',
},
{
text: 'How are you?',
},
],
},
]

const expected = [
{
role: 'user',
content: 'Hello, world!',
},
{
role: 'bot',
content: 'Hello, user!\nHow are you?',
},
]

expect(formatChat(messages)).toEqual(expected)
})

it('should throw an error if the content is missing', () => {
const messages = [
{
role: 'user',
},
]

expect(() => formatChat(messages)).toThrowError(
"Missing 'content' (ref: https://github.com/ggerganov/llama.cpp/issues/8367)",
)
})

it('should throw an error if the content type is invalid', () => {
const messages = [
{
role: 'user',
content: 42,
},
]

expect(() => formatChat(messages)).toThrowError(
"Invalid 'content' type (ref: https://github.com/ggerganov/llama.cpp/issues/8367)",
)
})
})
47 changes: 47 additions & 0 deletions src/chat.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
export type RNLlamaMessagePart = {
text?: string
}

export type RNLlamaOAICompatibleMessage = {
role: string
content?: string | RNLlamaMessagePart[] | any // any for check invalid content type
}

export type RNLlamaChatMessage = {
role: string
content: string
}

export function formatChat(
messages: RNLlamaOAICompatibleMessage[],
): RNLlamaChatMessage[] {
const chat: RNLlamaChatMessage[] = []

messages.forEach((currMsg) => {
const role: string = currMsg.role || ''

let content: string = ''
if ('content' in currMsg) {
if (typeof currMsg.content === 'string') {
;({ content } = currMsg)
} else if (Array.isArray(currMsg.content)) {
currMsg.content.forEach((part) => {
if ('text' in part) {
content += `${content ? '\n' : ''}${part.text}`
}
})
} else {
throw new TypeError(
"Invalid 'content' type (ref: https://github.com/ggerganov/llama.cpp/issues/8367)",
)
}
} else {
throw new Error(
"Missing 'content' (ref: https://github.com/ggerganov/llama.cpp/issues/8367)",
)
}

chat.push({ role, content })
})
return chat
}

0 comments on commit a75662c

Please sign in to comment.