react-native-litert-lm 0.3.7 → 0.4.0

package/android/src/main/java/com/margelo/nitro/dev/litert/litertlm/HybridLiteRTLM.kt

@@ -27,6 +27,10 @@ import com.margelo.nitro.dev.litert.litertlm.Role
 import com.margelo.nitro.core.Promise
 import com.google.ai.edge.litertlm.Content
 import com.google.ai.edge.litertlm.Contents
+import com.google.ai.edge.litertlm.ExperimentalApi
+import com.google.ai.edge.litertlm.ExperimentalFlags
+import com.google.ai.edge.litertlm.OpenApiTool
+import com.google.ai.edge.litertlm.ToolProvider
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicBoolean
@@ -167,6 +171,8 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     private var topP: Double = 0.95
     private var maxTokens: Int = 1024
     private var systemPrompt: String? = null
+    private var tools: Array<ToolDefinition>? = null
+    private var enableSpeculativeDecoding: Boolean = false
 
     override val memorySize: Long
         get() = 1024L * 1024L * 1024L // ~1GB (models are large)
@@ -196,8 +202,13 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                     cfg.topP?.let { topP = it }
                     cfg.maxTokens?.let { maxTokens = it.toInt() }
                     cfg.systemPrompt?.let { systemPrompt = it }
+                    tools = cfg.tools
+                    enableSpeculativeDecoding = cfg.enableSpeculativeDecoding ?: false
                 }
     
+                // Whether to run engine validation after loading
+                val shouldValidate = config?.validate?: false
+    
                 try {
                     // Early GPU hardware check: probe for OpenCL library before
                     // spending time on engine creation. LiteRT-LM's GPU delegate
@@ -238,12 +249,12 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                         else -> com.google.ai.edge.litertlm.Backend.CPU()
                     }
                     
-                    // Detect multimodal support from model filename.
+                    // Detect multimodal support. Check config.multimodal flag first, then fall back to filename sniffing.
                     // Only Gemma 3n bundles vision/audio executors; Gemma 4 E2B is text-only.
                     // Passing vision/audio backends to a text-only model causes
                     // vision_litert_compiled_model_executor init failures.
                     val modelFileName = modelPath.substringAfterLast("/").lowercase()
-                    val isMultimodal = modelFileName.contains("3n") || modelFileName.contains("gemma3")
+                    val isMultimodal = config?.multimodal ?: (modelFileName.contains("3n") || modelFileName.contains("gemma3"))
     
                     val lmVisionBackend = if (isMultimodal) com.google.ai.edge.litertlm.Backend.GPU() else null
                     val lmAudioBackend = if (isMultimodal) com.google.ai.edge.litertlm.Backend.CPU() else null
@@ -275,6 +286,11 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     
                     if (isClosed) return@synchronized
                     
+                    if (enableSpeculativeDecoding) {
+                        @OptIn(ExperimentalApi::class)
+                        ExperimentalFlags.enableSpeculativeDecoding = true
+                    }
+
                     // Initialize Engine
                     engine = Engine(engineConfig).also { it.initialize() }
                     Log.i(TAG, "Engine created and initialized successfully")
@@ -284,8 +300,17 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                     Log.i(TAG, "Conversation created successfully")
 
                     // Validate the engine actually works with a quick test inference.
-                    // GPU backend can initialize without error but silently fail to produce tokens.
-                    validateEngine()
+                    // GPU/NPU backends can initialize without error but silently fail to
+                    // produce tokens — enabling this catches those failures at load time.
+                    // CPU is always reliable so validation is never run on it, even when
+                    // the `validate` flag is set.
+                    if (shouldValidate) {
+                        if (backend == Backend.GPU || backend == Backend.NPU) {
+                            validateEngine()
+                        } else {
+                            Log.i(TAG, "Validation skipped: CPU backend is always reliable")
+                        }
+                    }
     
                 } catch (e: Exception) {
                     Log.e(TAG, "Failed to load model: ${e.message}", e)
@@ -349,32 +374,48 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     // -------------------------------------------------------------------------
     // sendMessageAsync - Streaming inference
     // -------------------------------------------------------------------------
-    override fun sendMessageAsync(message: String, onToken: (String, Boolean) -> Unit) {
-        // This is already async (void return), so we execute immediately on the calling thread
-        // (which is the Nitro specialized thread, not Main).
-        // The SDK's sendMessageAsync is non-blocking anyway.
-        ensureLoaded()
+    override fun sendMessageAsync(message: String, onToken: (String, Boolean) -> Unit): Promise<Unit> {
+        return Promise.parallel {
+            val latch = CountDownLatch(1)
+            val errorRef = AtomicReference<Throwable?>(null)
 
-        // Add user message to history
-        history.add(Message(Role.USER, message))
-        Log.d(TAG, "sendMessageAsync: $message")
+            ensureLoaded()
 
-        val fullResponseBuilder = StringBuilder()
-        
-        val listener = StreamingCallbackListener(
-            onToken = onToken,
-            responseBuilder = fullResponseBuilder,
-            history = history,
-            userMessage = message,
-            onStatsReady = { stats -> lastStats = stats },
-        )
+            // Add user message to history
+            history.add(Message(Role.USER, message))
+            Log.d(TAG, "sendMessageAsync: $message")
 
-        try {
-            val userMsg = LiteRTMessage.user(message)
-            conversation!!.sendMessageAsync(message = userMsg, callback = listener)
-        } catch (e: Exception) {
-            Log.e(TAG, "Failed to initiate async generation", e)
-            onToken("Error: ${e.message}", true)
+            val fullResponseBuilder = StringBuilder()
+            
+            val listener = StreamingCallbackListener(
+                onToken = { token, done ->
+                    onToken(token, done)
+                    if (done) {
+                        latch.countDown()
+                    }
+                },
+                responseBuilder = fullResponseBuilder,
+                history = history,
+                userMessage = message,
+                onStatsReady = { stats -> lastStats = stats },
+            )
+
+            try {
+                val userMsg = LiteRTMessage.user(message)
+                conversation!!.sendMessageAsync(message = userMsg, callback = listener)
+            } catch (e: Exception) {
+                Log.e(TAG, "Failed to initiate async generation", e)
+                errorRef.set(e)
+                onToken("Error: ${e.message}", true)
+                latch.countDown()
+            }
+
+            // Wait for completion or error
+            latch.await()
+            val err = errorRef.get()
+            if (err != null) {
+                throw RuntimeException("Async inference failed: ${err.message}", err)
+            }
         }
     }
 
@@ -464,6 +505,14 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
         return Promise.parallel {
             Log.i(TAG, "downloadModel: $url -> $fileName")
             
+            if (!url.startsWith("https://", ignoreCase = true)) {
+                throw IllegalArgumentException("Invalid download URL: HTTPS is required for security.")
+            }
+            
+            if (fileName.contains("..") || fileName.contains("/") || fileName.contains("\\")) {
+                throw IllegalArgumentException("Invalid filename: path traversal or directory separators are not allowed.")
+            }
+            
             val context = LiteRTLMInitProvider.applicationContext ?: throw RuntimeException("Context not available")
             val modelsDir = java.io.File(context.filesDir, "models")
             if (!modelsDir.exists()) {
@@ -545,6 +594,11 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     override fun deleteModel(fileName: String): Promise<Unit> {
         return Promise.parallel {
             Log.i(TAG, "deleteModel: $fileName")
+            
+            if (fileName.contains("..") || fileName.contains("/") || fileName.contains("\\")) {
+                throw IllegalArgumentException("Invalid filename: path traversal or directory separators are not allowed.")
+            }
+            
             val context = LiteRTLMInitProvider.applicationContext ?: throw RuntimeException("Context not available")
             val modelsDir = java.io.File(context.filesDir, "models")
             val modelFile = java.io.File(modelsDir, fileName)
@@ -676,19 +730,9 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
 
     private fun cleanupInternal() {
         try {
+            conversation?.close()
             conversation = null
-            // Explicitly close engine if it supports it to free native memory immediately
-            // Assuming Engine implements AutoCloseable or has close()
-            if (engine is AutoCloseable) {
-                (engine as AutoCloseable).close()
-            } else {
-                 // Try reflection or just null it if no close method
-                try {
-                    engine?.javaClass?.getMethod("close")?.invoke(engine)
-                } catch (e: Exception) {
-                    // Method not found, rely on GC
-                }
-            }
+            engine?.close()        // Direct call
             engine = null 
         } catch (e: Exception) {
             Log.e(TAG, "Error closing resources", e)
@@ -706,41 +750,37 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
         // v0.10.2 enforces single-session: close existing conversation first
         conversation?.let { oldConv ->
             try {
-                if (oldConv is AutoCloseable) {
-                    oldConv.close()
-                } else {
-                    oldConv.javaClass.getMethod("close").invoke(oldConv)
-                }
+                oldConv.close()
             } catch (e: Exception) {
                 Log.w(TAG, "Failed to close old conversation: ${e.message}")
             }
             conversation = null
         }
+        // Map tools
+        val lmTools: List<ToolProvider>? = tools?.map { tool ->
+            val apiTool = object : OpenApiTool {
+                override fun getToolDescriptionJsonString(): String {
+                    return tool.parametersJson
+                }
+                override fun execute(paramsJsonString: String): String {
+                    return "{}"
+                }
+            }
+            (apiTool as Any) as ToolProvider
+        }
+
         // Create conversation with explicit SamplerConfig (required by Gallery pattern).
         // GPU backend may fail silently without proper sampler params.
         val convConfig = ConversationConfig(
             samplerConfig = SamplerConfig(
                 topK = topK,
-                topP = topP,
-                temperature = temperature,
-            )
+                topP = topP.toDouble(),
+                temperature = temperature.toDouble(),
+            ),
+            systemInstruction = systemPrompt?.let { Contents.of(Content.Text(it)) },
+            tools = lmTools ?: emptyList()
         )
         conversation = engine!!.createConversation(convConfig)
-        // Apply system prompt/instruction if set
-        systemPrompt?.let { prompt ->
-            if (prompt.isNotEmpty()) {
-                try {
-                    // Send system instruction as the first turn to prime the conversation.
-                    // LiteRT-LM's Conversation API handles chat template formatting,
-                    // including Gemma's <start_of_turn>system block.
-                    val systemMsg = LiteRTMessage.system(prompt)
-                    conversation!!.sendMessage(message = systemMsg)
-                    Log.i(TAG, "System prompt applied (${prompt.length} chars)")
-                } catch (e: Exception) {
-                    Log.w(TAG, "Failed to apply system prompt: ${e.message}")
-                }
-            }
-        }
     }
 
     /**
@@ -815,5 +855,62 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
         createNewConversation()
     }
 
+    override fun sendMultimodalMessage(parts: Array<MultimodalPart>): Promise<String> {
+        return Promise.parallel {
+            ensureLoaded()
+            val contents = mutableListOf<Content>()
+            var userTextRepresentation = ""
+            for (part in parts) {
+                when (part.type) {
+                    PartType.TEXT -> part.text?.let { 
+                        contents.add(Content.Text(it))
+                        userTextRepresentation += "$it "
+                    }
+                    PartType.IMAGE -> part.imageBuffer?.let { buffer ->
+                        val byteBuffer = buffer.getBuffer(false)
+                        val bytes = ByteArray(byteBuffer.remaining())
+                        byteBuffer.get(bytes)
+                        contents.add(Content.ImageBytes(bytes))
+                        userTextRepresentation += "[Image Buffer] "
+                    }
+                    PartType.AUDIO -> part.audioBuffer?.let { buffer ->
+                        val byteBuffer = buffer.getBuffer(false)
+                        val bytes = ByteArray(byteBuffer.remaining())
+                        byteBuffer.get(bytes)
+                        contents.add(Content.AudioBytes(bytes))
+                        userTextRepresentation += "[Audio Buffer] "
+                    }
+                }
+            }
+            userTextRepresentation = userTextRepresentation.trim()
+            history.add(Message(Role.USER, userTextRepresentation))
+
+            val userMsg = LiteRTMessage.user(Contents.of(contents))
+            val startTime = System.nanoTime()
+            val responseMsg = conversation!!.sendMessage(message = userMsg)
+            val elapsedMs = (System.nanoTime() - startTime) / 1_000_000.0
+
+            val response = responseMsg.contents.contents
+                .filterIsInstance<Content.Text>()
+                .joinToString("") { it.text }
+
+            history.add(Message(Role.MODEL, response))
+            
+            val promptTokens = userTextRepresentation.length / 4.0
+            val completionTokens = response.length / 4.0
+            lastStats = GenerationStats(
+                promptTokens = promptTokens,
+                completionTokens = completionTokens,
+                totalTokens = promptTokens + completionTokens,
+                timeToFirstToken = 0.0,
+                totalTime = elapsedMs,
+                tokensPerSecond = if (elapsedMs > 0) completionTokens / (elapsedMs / 1000.0) else 0.0
+            )
+            response
+        }
+    }
 
+    override fun countTokens(text: String): Double {
+        return -1.0
+    }
 }

package/android/src/main/java/dev/litert/litertlm/LiteRTLMPackage.kt

@@ -1,18 +1,35 @@
 package dev.litert.litertlm
 
+import android.os.Build
+import android.util.Log
 import com.facebook.react.TurboReactPackage
 import com.facebook.react.bridge.NativeModule
 import com.facebook.react.bridge.ReactApplicationContext
 import com.facebook.react.module.model.ReactModuleInfo
 import com.facebook.react.module.model.ReactModuleInfoProvider
-import com.margelo.nitro.core.HybridObject
 
 
 import com.margelo.nitro.dev.litert.litertlm.LiteRTLMOnLoad
 
 class LiteRTLMPackage : TurboReactPackage() {
+    companion object {
+        private const val TAG = "LiteRTLMPackage"
+
+        private fun isSupportedPrimaryAbi(): Boolean {
+            val primaryAbi = Build.SUPPORTED_64_BIT_ABIS.firstOrNull() ?: return false
+            return primaryAbi == "arm64-v8a"
+        }
+    }
     init {
-        LiteRTLMOnLoad.initializeNative()
+        if (!isSupportedPrimaryAbi()) {
+            Log.w(TAG, "Skipping LiteRTLM native init on unsupported primary ABI: ${Build.SUPPORTED_64_BIT_ABIS.firstOrNull()}")
+        } else {
+            try {
+                LiteRTLMOnLoad.initializeNative()
+            } catch (e: UnsatisfiedLinkError) {
+                Log.e(TAG, "LiteRTLM native init failed; disabling LiteRTLM for this process.", e)
+            }
+        }
     }
 
 

package/android/src/test/java/com/margelo/nitro/core/Promise.kt

@@ -0,0 +1,46 @@
+package com.margelo.nitro.core
+
+import androidx.annotation.Keep
+import com.facebook.proguard.annotations.DoNotStrip
+
+@Keep
+@DoNotStrip
+class Promise<T> {
+    companion object {
+        @JvmStatic
+        fun <T> parallel(block: () -> T): Promise<T> {
+            val promise = Promise<T>()
+            try {
+                val result = block()
+                promise.resolve(result)
+            } catch (e: Throwable) {
+                promise.reject(e)
+            }
+            return promise
+        }
+    }
+
+    var result: T? = null
+        private set
+    var error: Throwable? = null
+        private set
+    var isCompleted = false
+        private set
+    private val callbacks = mutableListOf<(T?, Throwable?) -> Unit>()
+
+    fun resolve(value: T) {
+        synchronized(this) {
+            result = value
+            isCompleted = true
+            callbacks.forEach { it(value, null) }
+        }
+    }
+
+    fun reject(exception: Throwable) {
+        synchronized(this) {
+            error = exception
+            isCompleted = true
+            callbacks.forEach { it(null, exception) }
+        }
+    }
+}

package/android/src/test/java/com/margelo/nitro/dev/litert/litertlm/HybridLiteRTLMTest.kt

@@ -0,0 +1,83 @@
+package com.margelo.nitro.dev.litert.litertlm
+
+import org.junit.Assert.*
+import org.junit.Before
+import org.junit.After
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.RuntimeEnvironment
+import dev.litert.litertlm.LiteRTLMInitProvider
+import java.lang.IllegalArgumentException
+
+@RunWith(RobolectricTestRunner::class)
+class HybridLiteRTLMTest {
+    private lateinit var bridge: HybridLiteRTLM
+
+    @Before
+    fun setUp() {
+        // Initialize the static applicationContext inside LiteRTLMInitProvider via reflection
+        try {
+            val field = LiteRTLMInitProvider::class.java.getDeclaredField("applicationContext")
+            field.isAccessible = true
+            field.set(null, RuntimeEnvironment.getApplication())
+        } catch (e: Exception) {
+            e.printStackTrace()
+        }
+        
+        bridge = HybridLiteRTLM()
+    }
+
+    @After
+    fun tearDown() {
+        bridge.close()
+    }
+
+    @Test
+    fun testAndroidPathTraversalPrevention() {
+        val traversals = arrayOf("../secret", "/etc/hosts", "nested\\..\\file", "..", "../", "..\\")
+        for (traversal in traversals) {
+            val promise = bridge.deleteModel(traversal)
+            assertNotNull("Promise should not be null", promise)
+            assertTrue("Promise should be completed", promise.isCompleted)
+            assertNotNull("Promise should have rejected with an error for filename: $traversal", promise.error)
+            val error = promise.error!!
+            val errMsg = error.message ?: error.cause?.message ?: ""
+            assertTrue("Expected message to contain traversal warning, got: $errMsg",
+                errMsg.contains("path traversal or directory separators are not allowed"))
+        }
+    }
+
+    @Test
+    fun testAndroidHTTPSDownloadEnforcement() {
+        val promise = bridge.downloadModel("http://insecure.site/model.bin", "model.bin", null)
+        assertNotNull("Promise should not be null", promise)
+        assertTrue("Promise should be completed", promise.isCompleted)
+        assertNotNull("Promise should have rejected with an error", promise.error)
+        val error = promise.error!!
+        val errMsg = error.message ?: error.cause?.message ?: ""
+        assertTrue("Expected message to contain HTTPS warning, got: $errMsg",
+            errMsg.contains("HTTPS is required for security"))
+    }
+
+    @Test
+    fun testAndroidMemoryTelemetry() {
+        val mem = bridge.getMemoryUsage()
+        assertNotNull(mem)
+        assertTrue(mem.nativeHeapBytes >= 0.0)
+        assertTrue(mem.residentBytes >= 0.0)
+        assertTrue(mem.availableMemoryBytes >= 0.0)
+    }
+
+    @Test
+    fun testAndroidInitialStats() {
+        val stats = bridge.getStats()
+        assertNotNull(stats)
+        assertEquals(0.0, stats.promptTokens, 0.0)
+        assertEquals(0.0, stats.completionTokens, 0.0)
+        assertEquals(0.0, stats.totalTokens, 0.0)
+        assertEquals(0.0, stats.timeToFirstToken, 0.0)
+        assertEquals(0.0, stats.totalTime, 0.0)
+        assertEquals(0.0, stats.tokensPerSecond, 0.0)
+    }
+}