bugfix: Fix memory clear bug

it5prasoon · it5prasoon · commit 289eb6ca6f16 · 2025-12-20T09:42:44.000+05:30
diff --git a/app/src/main/cpp/llama_context_wrapper.cpp b/app/src/main/cpp/llama_context_wrapper.cpp
@@ -194,8 +194,18 @@ void LlamaContextWrapper::generateStream(const std::string& prompt, TokenCallbac
         return;
     }
     
-    // KV cache is already empty for a new generation
-    // No need to clear explicitly
+    // Clear memory (KV cache) for new generation - new API uses llama_memory_clear
+    llama_memory_t mem = llama_get_memory(context_);
+    if (mem != nullptr) {
+        llama_memory_clear(mem, true);
+        LOGD("Memory cleared for new generation");
+    }
+    
+    // Reset sampler state for new generation
+    if (sampler_ != nullptr) {
+        llama_sampler_reset(sampler_);
+        LOGD("Sampler reset for new generation");
+    }
     
     // Create batch for prompt processing
     llama_batch batch = llama_batch_init(cfg.batchSize, 0, 1);
diff --git a/app/src/main/cpp/llama_jni.cpp b/app/src/main/cpp/llama_jni.cpp
@@ -45,12 +45,47 @@ static jstring stringToJstring(JNIEnv* env, const std::string& str) {
     return env->NewStringUTF(str.c_str());
 }
 
-// Helper to throw Java exception
+// Helper to throw Java exception - uses RuntimeException for safety
 static void throwException(JNIEnv* env, const char* className, const char* message) {
+    // Try the specified class first
     jclass exClass = env->FindClass(className);
-    if (exClass != nullptr) {
+    if (exClass != nullptr && !env->ExceptionCheck()) {
         env->ThrowNew(exClass, message);
         env->DeleteLocalRef(exClass);
+    } else {
+        // Fall back to RuntimeException if class not found
+        env->ExceptionClear();
+        jclass runtimeEx = env->FindClass("java/lang/RuntimeException");
+        if (runtimeEx != nullptr) {
+            env->ThrowNew(runtimeEx, message);
+            env->DeleteLocalRef(runtimeEx);
+        }
+    }
+}
+
+// Helper to throw LlamaException.GenerationError
+static void throwGenerationError(JNIEnv* env, const char* message) {
+    jclass exClass = env->FindClass("com/llamakotlin/android/exception/LlamaException$GenerationError");
+    if (exClass != nullptr && !env->ExceptionCheck()) {
+        // Find constructor
+        jmethodID constructor = env->GetMethodID(exClass, "<init>", "(Ljava/lang/String;Ljava/lang/Throwable;)V");
+        if (constructor != nullptr) {
+            jstring jmsg = env->NewStringUTF(message);
+            jobject ex = env->NewObject(exClass, constructor, jmsg, nullptr);
+            if (ex != nullptr) {
+                env->Throw((jthrowable)ex);
+                env->DeleteLocalRef(ex);
+            }
+            env->DeleteLocalRef(jmsg);
+        } else {
+            // Fallback to RuntimeException
+            env->ExceptionClear();
+            throwException(env, "java/lang/RuntimeException", message);
+        }
+        env->DeleteLocalRef(exClass);
+    } else {
+        env->ExceptionClear();
+        throwException(env, "java/lang/RuntimeException", message);
     }
 }
 
@@ -176,7 +211,7 @@ Java_com_llamakotlin_android_LlamaNative_nativeLoadModel(
     
     if (!success) {
         std::string error = context->getLastError();
-        throwException(env, "com/llamakotlin/android/exception/LlamaException", error.c_str());
+        throwGenerationError(env, error.c_str());
         return JNI_FALSE;
     }
     
@@ -242,8 +277,7 @@ Java_com_llamakotlin_android_LlamaNative_nativeGenerate(
     std::string result = context->generate(promptStr, configPtr);
     
     if (result.empty() && !context->getLastError().empty()) {
-        throwException(env, "com/llamakotlin/android/exception/LlamaException", 
-                       context->getLastError().c_str());
+        throwGenerationError(env, context->getLastError().c_str());
         return nullptr;
     }
     
@@ -318,10 +352,11 @@ Java_com_llamakotlin_android_LlamaNative_nativeGenerateStream(
     env->DeleteGlobalRef(globalCallback);
     env->DeleteLocalRef(callbackClass);
     
-    // Check for errors
-    if (!context->getLastError().empty()) {
-        throwException(env, "com/llamakotlin/android/exception/LlamaException",
-                       context->getLastError().c_str());
+    // Check for errors - don't throw if already completed successfully
+    std::string error = context->getLastError();
+    if (!error.empty()) {
+        LOGE("Generation error: %s", error.c_str());
+        throwGenerationError(env, error.c_str());
     }
 }
 
diff --git a/sample/src/main/java/com/llamakotlin/sample/MainActivity.kt b/sample/src/main/java/com/llamakotlin/sample/MainActivity.kt
@@ -1,12 +1,15 @@
 package com.llamakotlin.sample
 
+import android.content.Context
+import android.content.Intent
+import android.content.SharedPreferences
+import android.net.Uri
 import android.os.Bundle
-import android.os.Environment
+import android.provider.OpenableColumns
 import android.widget.Toast
 import androidx.activity.result.contract.ActivityResultContracts
 import androidx.appcompat.app.AppCompatActivity
 import androidx.lifecycle.lifecycleScope
-import com.llamakotlin.android.LlamaConfig
 import com.llamakotlin.android.LlamaModel
 import com.llamakotlin.android.exception.LlamaException
 import com.llamakotlin.sample.databinding.ActivityMainBinding
@@ -16,20 +19,52 @@ import kotlinx.coroutines.flow.catch
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import java.io.File
+import java.io.FileOutputStream
 
 class MainActivity : AppCompatActivity() {
 
     private lateinit var binding: ActivityMainBinding
+    private lateinit var prefs: SharedPreferences
     private var model: LlamaModel? = null
     private var generationJob: Job? = null
 
+    companion object {
+        private const val PREFS_NAME = "llama_prefs"
+        private const val KEY_MODEL_PATH = "model_path"
+    }
+
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
-        binding = ActivityMainBinding.inflate(layoutInflater)
-        setContentView(binding.root)
+        try {
+            binding = ActivityMainBinding.inflate(layoutInflater)
+            setContentView(binding.root)
+            
+            prefs = getSharedPreferences(PREFS_NAME, Context.MODE_PRIVATE)
 
-        setupUI()
-        showVersion()
+            setupUI()
+            showVersion()
+            
+            // Auto-load last used model
+            autoLoadModel()
+        } catch (e: Exception) {
+            e.printStackTrace()
+            Toast.makeText(this, "Error: ${e.message}", Toast.LENGTH_LONG).show()
+        }
+    }
+    
+    private fun autoLoadModel() {
+        val savedPath = prefs.getString(KEY_MODEL_PATH, null)
+        if (savedPath != null) {
+            val modelFile = File(savedPath)
+            if (modelFile.exists()) {
+                binding.tvStatus.text = "Auto-loading last model..."
+                loadModel(savedPath)
+            } else {
+                // Clear saved path if file no longer exists
+                prefs.edit().remove(KEY_MODEL_PATH).apply()
+                binding.tvStatus.text = "Previous model not found. Please select a new model."
+            }
+        }
     }
 
     private fun setupUI() {
@@ -63,39 +98,94 @@ class MainActivity : AppCompatActivity() {
         }
     }
 
-    private val pickFileLauncher = registerForActivityResult(
+    // File picker for GGUF models
+    private val filePickerLauncher = registerForActivityResult(
         ActivityResultContracts.OpenDocument()
-    ) { uri ->
-        uri?.let {
-            // Copy to local storage or use content resolver
-            val path = getModelPath(uri.path ?: "")
-            if (path.isNotEmpty()) {
-                loadModel(path)
-            } else {
-                Toast.makeText(this, "Invalid model file", Toast.LENGTH_SHORT).show()
-            }
+    ) { uri: Uri? ->
+        uri?.let { selectedUri ->
+            // Take persistent permission
+            contentResolver.takePersistableUriPermission(
+                selectedUri,
+                Intent.FLAG_GRANT_READ_URI_PERMISSION
+            )
+            copyAndLoadModel(selectedUri)
         }
     }
 
     private fun pickModelFile() {
-        // For simplicity, look for model in Downloads folder
-        val downloadsDir = Environment.getExternalStoragePublicDirectory(Environment.DIRECTORY_DOWNLOADS)
-        val ggufFiles = downloadsDir.listFiles { file -> file.extension == "gguf" }
+        binding.tvStatus.text = "Opening file picker...\nSelect a .gguf model file"
         
-        if (ggufFiles.isNullOrEmpty()) {
-            Toast.makeText(this, "No .gguf files in Downloads. Please download a model first.", Toast.LENGTH_LONG).show()
-            binding.tvStatus.text = "No model found. Download a .gguf model to Downloads folder."
-            return
-        }
+        // Use Storage Access Framework - works on all Android versions
+        filePickerLauncher.launch(arrayOf("*/*"))
+    }
+
+    private fun copyAndLoadModel(uri: Uri) {
+        binding.tvStatus.text = "Copying model to app storage..."
+        binding.btnLoadModel.isEnabled = false
+
+        lifecycleScope.launch(Dispatchers.IO) {
+            try {
+                // Get file name
+                val fileName = getFileName(uri) ?: "model.gguf"
+                
+                if (!fileName.endsWith(".gguf", ignoreCase = true)) {
+                    withContext(Dispatchers.Main) {
+                        binding.tvStatus.text = "Please select a .gguf file"
+                        binding.btnLoadModel.isEnabled = true
+                    }
+                    return@launch
+                }
+
+                // Copy to app's internal storage
+                val modelFile = File(filesDir, fileName)
+                
+                withContext(Dispatchers.Main) {
+                    binding.tvStatus.text = "Copying $fileName...\nThis may take a few minutes for large models."
+                }
+
+                contentResolver.openInputStream(uri)?.use { input ->
+                    FileOutputStream(modelFile).use { output ->
+                        val buffer = ByteArray(8192)
+                        var bytesRead: Int
+                        var totalBytes = 0L
+                        while (input.read(buffer).also { bytesRead = it } != -1) {
+                            output.write(buffer, 0, bytesRead)
+                            totalBytes += bytesRead
+                            // Update progress every 10MB
+                            if (totalBytes % (10 * 1024 * 1024) < 8192) {
+                                withContext(Dispatchers.Main) {
+                                    binding.tvStatus.text = "Copying $fileName...\n${totalBytes / (1024 * 1024)} MB copied"
+                                }
+                            }
+                        }
+                    }
+                }
 
-        // Use the first found model
-        val modelFile = ggufFiles.first()
-        loadModel(modelFile.absolutePath)
+                withContext(Dispatchers.Main) {
+                    binding.tvStatus.text = "Model copied. Loading..."
+                    binding.btnLoadModel.isEnabled = true
+                    loadModel(modelFile.absolutePath)
+                }
+
+            } catch (e: Exception) {
+                e.printStackTrace()
+                withContext(Dispatchers.Main) {
+                    binding.tvStatus.text = "Error copying model: ${e.message}"
+                    binding.btnLoadModel.isEnabled = true
+                }
+            }
+        }
     }
 
-    private fun getModelPath(uriPath: String): String {
-        // Simple path extraction - in production use ContentResolver
-        return uriPath
+    private fun getFileName(uri: Uri): String? {
+        var name: String? = null
+        contentResolver.query(uri, null, null, null, null)?.use { cursor ->
+            val nameIndex = cursor.getColumnIndex(OpenableColumns.DISPLAY_NAME)
+            if (cursor.moveToFirst() && nameIndex >= 0) {
+                name = cursor.getString(nameIndex)
+            }
+        }
+        return name ?: uri.lastPathSegment
     }
 
     private fun loadModel(modelPath: String) {
@@ -114,6 +204,9 @@ class MainActivity : AppCompatActivity() {
                     topK = 40
                 }
                 
+                // Save model path for auto-loading
+                prefs.edit().putString(KEY_MODEL_PATH, modelPath).apply()
+                
                 binding.tvStatus.text = "Model loaded: ${File(modelPath).name}"
                 binding.btnGenerate.isEnabled = true
                 Toast.makeText(this@MainActivity, "Model loaded successfully!", Toast.LENGTH_SHORT).show()
diff --git a/sample/src/main/res/layout/activity_main.xml b/sample/src/main/res/layout/activity_main.xml
@@ -3,13 +3,15 @@
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     android:orientation="vertical"
-    android:padding="16dp">
+    android:padding="16dp"
+    android:fitsSystemWindows="true">
 
     <!-- Status -->
     <TextView
         android:id="@+id/tvStatus"
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
+        android:layout_marginTop="?attr/actionBarSize"
         android:text="No model loaded"
         android:textColor="?android:attr/textColorSecondary"
         android:padding="8dp"
@@ -21,7 +23,7 @@
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
         android:layout_marginTop="8dp"
-        android:text="Load Model from Downloads" />
+        android:text="Select Model File (.gguf)" />
 
     <!-- Prompt Input -->
     <EditText
diff --git a/sample/src/main/res/mipmap-anydpi-v26/ic_launcher.xml b/sample/src/main/res/mipmap-anydpi-v26/ic_launcher.xml
@@ -0,0 +1,5 @@
+<?xml version="1.0" encoding="utf-8"?>
+<adaptive-icon xmlns:android="http://schemas.android.com/apk/res/android">
+    <background android:drawable="@color/purple_500"/>
+    <foreground android:drawable="@color/white"/>
+</adaptive-icon>
diff --git a/sample/src/main/res/mipmap-anydpi-v26/ic_launcher_round.xml b/sample/src/main/res/mipmap-anydpi-v26/ic_launcher_round.xml
@@ -0,0 +1,5 @@
+<?xml version="1.0" encoding="utf-8"?>
+<adaptive-icon xmlns:android="http://schemas.android.com/apk/res/android">
+    <background android:drawable="@color/purple_500"/>
+    <foreground android:drawable="@color/white"/>
+</adaptive-icon>
