cui-llama.rn 1.0.3 → 1.0.6

package/src/NativeRNLlama.ts

@@ -19,7 +19,6 @@ export type NativeContextParams = {
 
   lora?: string // lora_adaptor
   lora_scaled?: number
-  lora_base?: string
 
   rope_freq_base?: number
   rope_freq_scale?: number
@@ -112,6 +111,11 @@ export type NativeSessionLoadResult = {
   prompt: string
 }
 
+export type NativeLlamaChatMessage = {
+  role: string
+  content: string
+}
+
 export interface Spec extends TurboModule {
   setContextLimit(limit: number): Promise<void>
   initContext(params: NativeContextParams): Promise<NativeLlamaContext>
@@ -132,6 +136,11 @@ export interface Spec extends TurboModule {
   stopCompletion(contextId: number): Promise<void>
   tokenizeAsync(contextId: number, text: string): Promise<NativeTokenizeResult>
   tokenizeSync(contextId: number, text: string): NativeTokenizeResult
+  getFormattedChat(
+    contextId: number,
+    messages: NativeLlamaChatMessage[],
+    chatTemplate?: string,
+  ): Promise<string>
   detokenize(contextId: number, tokens: number[]): Promise<string>
   embedding(contextId: number, text: string): Promise<NativeEmbeddingResult>
   bench(

package/src/chat.ts

@@ -0,0 +1,44 @@
+import type { NativeLlamaChatMessage } from './NativeRNLlama'
+
+export type RNLlamaMessagePart = {
+  text?: string
+}
+
+export type RNLlamaOAICompatibleMessage = {
+  role: string
+  content?: string | RNLlamaMessagePart[] | any // any for check invalid content type
+}
+
+export function formatChat(
+  messages: RNLlamaOAICompatibleMessage[],
+): NativeLlamaChatMessage[] {
+  const chat: NativeLlamaChatMessage[] = []
+
+  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
+}

package/src/index.ts

@@ -12,6 +12,8 @@ import type {
   NativeSessionLoadResult,
 } from './NativeRNLlama'
 import { SchemaGrammarConverter, convertJsonSchemaToGrammar } from './grammar'
+import type { RNLlamaOAICompatibleMessage } from './chat'
+import { formatChat } from './chat'
 
 export { SchemaGrammarConverter, convertJsonSchemaToGrammar }
 
@@ -40,8 +42,11 @@ export type ContextParams = NativeContextParams
 
 export type CompletionParams = Omit<
   NativeCompletionParams,
-  'emit_partial_completion'
->
+  'emit_partial_completion' | 'prompt'
+> & {
+  prompt?: string
+  messages?: RNLlamaOAICompatibleMessage[]
+}
 
 export type BenchResult = {
   modelDesc: string
@@ -60,7 +65,9 @@ export class LlamaContext {
 
   reasonNoGPU: string = ''
 
-  model: Object = {}
+  model: {
+    isChatTemplateSupported?: boolean
+  } = {}
 
   constructor({ contextId, gpu, reasonNoGPU, model }: NativeLlamaContext) {
     this.id = contextId
@@ -74,7 +81,7 @@ export class LlamaContext {
    */
   async loadSession(filepath: string): Promise<NativeSessionLoadResult> {
     let path = filepath
-    if (filepath.startsWith(`file://`)) path = path.slice(7)
+    if (path.startsWith('file://')) path = path.slice(7)
     return RNLlama.loadSession(this.id, path)
   }
 
@@ -88,10 +95,27 @@ export class LlamaContext {
     return RNLlama.saveSession(this.id, filepath, options?.tokenSize || -1)
   }
 
+  async getFormattedChat(
+    messages: RNLlamaOAICompatibleMessage[],
+  ): Promise<string> {
+    const chat = formatChat(messages)
+    return RNLlama.getFormattedChat(
+      this.id,
+      chat,
+      this.model?.isChatTemplateSupported ? undefined : 'chatml',
+    )
+  }
+
   async completion(
     params: CompletionParams,
     callback?: (data: TokenData) => void,
   ): Promise<NativeCompletionResult> {
+
+    let finalPrompt = params.prompt
+    if (params.messages) { // messages always win
+      finalPrompt = await this.getFormattedChat(params.messages)
+    }
+
     let tokenListener: any =
       callback &&
       EventEmitter.addListener(EVENT_ON_TOKEN, (evt: TokenNativeEvent) => {
@@ -99,8 +123,11 @@ export class LlamaContext {
         if (contextId !== this.id) return
         callback(tokenResult)
       })
+
+    if (!finalPrompt) throw new Error('Prompt is required')
     const promise = RNLlama.completion(this.id, {
       ...params,
+      prompt: finalPrompt,
       emit_partial_completion: !!callback,
     })
     return promise