expo-ai-kit 0.3.5 → 0.3.6

package/android/src/main/java/expo/modules/aikit/ExpoAiKitModule.kt

@@ -49,18 +49,26 @@ class ExpoAiKitModule : Module() {
 
       // Build conversation history prompt from all non-system messages
       // On-device models are stateless, so we must include full history in each request
-      val conversationPrompt = messages
-        .filter { it["role"] != "system" }
-        .joinToString("\n") { msg ->
-          val role = (msg["role"] as? String ?: "user").uppercase()
-          val content = msg["content"] as? String ?: ""
-          "$role: $content"
-        } + "\nASSISTANT:"
+      val nonSystemMessages = messages.filter { it["role"] != "system" }
 
       // Route to active model
       val text = if (activeModelId == "mlkit") {
+        // ML Kit: use role-prefixed format since it has no conversation API
+        val conversationPrompt = nonSystemMessages
+          .joinToString("\n") { msg ->
+            val role = (msg["role"] as? String ?: "user").uppercase()
+            val content = msg["content"] as? String ?: ""
+            "$role: $content"
+          } + "\nASSISTANT:"
         promptClient.generateText(conversationPrompt, systemPrompt)
       } else {
+        // Gemma/LiteRT-LM: pass raw content — the Conversation API handles
+        // turn formatting internally. Adding "USER:"/"ASSISTANT:" markers
+        // causes double-formatting and garbled output.
+        val conversationPrompt = nonSystemMessages
+          .joinToString("\n") { msg ->
+            msg["content"] as? String ?: ""
+          }
         gemmaClient.generateText(conversationPrompt, systemPrompt)
       }
       mapOf("text" to text)
@@ -73,25 +81,15 @@ class ExpoAiKitModule : Module() {
         ?.get("content") as? String
         ?: fallbackSystemPrompt.ifBlank { "You are a helpful, friendly assistant." }
 
-      // Build conversation history prompt from all non-system messages
-      // On-device models are stateless, so we must include full history in each request
-      val conversationPrompt = messages
-        .filter { it["role"] != "system" }
-        .joinToString("\n") { msg ->
-          val role = (msg["role"] as? String ?: "user").uppercase()
-          val content = msg["content"] as? String ?: ""
-          "$role: $content"
-        } + "\nASSISTANT:"
+      val nonSystemMessages = messages.filter { it["role"] != "system" }
 
       // Launch streaming in a coroutine that can be cancelled
       val job = streamScope.launch {
-        val streamAccumulator = StringBuilder()
-        val streamCallback = { token: String, _: String, isDone: Boolean ->
-          streamAccumulator.append(token)
+        val streamCallback = { token: String, accumulatedText: String, isDone: Boolean ->
           sendEvent("onStreamToken", mapOf(
             "sessionId" to sessionId,
             "token" to token,
-            "accumulatedText" to streamAccumulator.toString(),
+            "accumulatedText" to accumulatedText,
             "isDone" to isDone
           ))
 
@@ -102,8 +100,20 @@ class ExpoAiKitModule : Module() {
 
         // Route to active model
         if (activeModelId == "mlkit") {
+          // ML Kit: use role-prefixed format since it has no conversation API
+          val conversationPrompt = nonSystemMessages
+            .joinToString("\n") { msg ->
+              val role = (msg["role"] as? String ?: "user").uppercase()
+              val content = msg["content"] as? String ?: ""
+              "$role: $content"
+            } + "\nASSISTANT:"
           promptClient.generateTextStream(conversationPrompt, systemPrompt, streamCallback)
         } else {
+          // Gemma/LiteRT-LM: pass raw content — Conversation API handles turn formatting
+          val conversationPrompt = nonSystemMessages
+            .joinToString("\n") { msg ->
+              msg["content"] as? String ?: ""
+            }
           gemmaClient.generateTextStream(conversationPrompt, systemPrompt, streamCallback)
         }
       }

package/android/src/main/java/expo/modules/aikit/GemmaInferenceClient.kt

@@ -219,22 +219,37 @@ class GemmaInferenceClient(private val context: Context) {
     try {
       withContext(Dispatchers.IO) {
         suspendCancellableCoroutine<Unit> { continuation ->
+          val accumulatedBuilder = StringBuilder()
           var previousText = ""
           conv.sendMessageAsync(
             Contents.of(fullPrompt),
             object : MessageCallback {
               override fun onMessage(message: Message) {
-                val accumulated = message.toString()
-                val token = if (accumulated.length > previousText.length) {
-                  accumulated.substring(previousText.length)
+                val messageText = message.toString()
+
+                // LiteRT-LM may deliver accumulated text or delta tokens depending
+                // on the version. Detect which by checking if messageText extends
+                // what we've seen before.
+                val token: String
+                if (messageText.startsWith(previousText) && messageText.length >= previousText.length) {
+                  // Accumulated text — extract delta
+                  token = messageText.substring(previousText.length)
+                  previousText = messageText
+                  accumulatedBuilder.clear()
+                  accumulatedBuilder.append(messageText)
                 } else {
-                  ""
+                  // Delta token — accumulate ourselves
+                  token = messageText
+                  accumulatedBuilder.append(messageText)
+                  previousText = accumulatedBuilder.toString()
                 }
-                previousText = accumulated
+
+                val accumulated = accumulatedBuilder.toString()
                 onChunk(token, accumulated, false)
               }
               override fun onDone() {
-                onChunk("", previousText, true)
+                val finalText = accumulatedBuilder.toString()
+                onChunk("", finalText, true)
                 continuation.resume(Unit)
               }
               override fun onError(throwable: Throwable) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "expo-ai-kit",
-  "version": "0.3.5",
+  "version": "0.3.6",
   "description": "On-device AI for Expo apps — run Gemma 4, Apple Foundation Models, and ML Kit locally with zero API keys",
   "main": "build/index.js",
   "types": "build/index.d.ts",