@capgo/capacitor-llm 8.0.19 → 8.1.0

package/android/src/main/java/ee/forgr/capgo_llm/LLM.java

@@ -1,40 +1,72 @@
 package ee.forgr.capgo_llm;
 
 import android.content.Context;
+import android.os.Build;
+import com.google.ai.edge.litertlm.Backend;
+import com.google.ai.edge.litertlm.Content;
+import com.google.ai.edge.litertlm.Conversation;
+import com.google.ai.edge.litertlm.ConversationConfig;
+import com.google.ai.edge.litertlm.Engine;
+import com.google.ai.edge.litertlm.EngineConfig;
+import com.google.ai.edge.litertlm.Message;
+import com.google.ai.edge.litertlm.SamplerConfig;
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.mediapipe.tasks.genai.llminference.LlmInference;
 import com.google.mediapipe.tasks.genai.llminference.LlmInference.LlmInferenceOptions;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Locale;
+import com.google.mediapipe.tasks.genai.llminference.LlmInferenceSession;
+import com.google.mlkit.genai.common.FeatureStatus;
+import com.google.mlkit.genai.common.StreamingCallback;
+import com.google.mlkit.genai.prompt.Candidate;
+import com.google.mlkit.genai.prompt.GenerateContentRequest;
+import com.google.mlkit.genai.prompt.GenerateContentResponse;
+import com.google.mlkit.genai.prompt.Generation;
+import com.google.mlkit.genai.prompt.TextPart;
+import com.google.mlkit.genai.prompt.java.GenerativeModelFutures;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.InputStream;
+import java.util.Collections;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
-import org.pytorch.executorch.extension.llm.LlmCallback;
-import org.pytorch.executorch.extension.llm.LlmModule;
 
 public class LLM {
 
-    private enum Engine {
+    private enum BackendType {
+        LITERT,
         MEDIAPIPE,
-        EXECUTORCH
+        GEMINI_NANO
     }
 
+    private static final double DEFAULT_TOP_P = 0.95d;
     private static LLM instance;
-    private LlmInference mediaPipeInference;
-    private LlmModule executorchModule;
-    private Map<String, ChatSession> chatSessions;
+
+    private final Context context;
+    private final Executor executor;
+    private final Map<String, ChatSession> chatSessions;
+
+    private Engine engine;
+    private GenerativeModelFutures geminiNanoModel;
+    private LlmInference llmInference;
     private boolean isReady = false;
-    private Context context;
-    private Executor executor;
     private String modelPath = null;
-    private String tokenizerPath = null;
-    private Engine currentEngine = Engine.MEDIAPIPE;
+    private String modelType = null;
+    private BackendType activeBackend = BackendType.LITERT;
+    private long currentLoadId = 0L;
+
+    private Integer maxTokens = 2048;
+    private Integer topk = 40;
+    private Float temperature = 0.1f;
+    private Integer randomSeed = 101;
 
     private LLM(Context context) {
         this.context = context;
-        this.chatSessions = new HashMap<>();
+        this.chatSessions = new ConcurrentHashMap<>();
         this.executor = Executors.newSingleThreadExecutor();
     }
 
@@ -45,109 +77,52 @@ public class LLM {
         return instance;
     }
 
-    // Model parameters
-    private Integer maxTokens = 2048;
-    private Integer sequenceLength = 2048;
-    private Integer topk = 40;
-    private Float temperature = 0.1f;
-    private String modelType = null;
-
     public void setModel(
         String path,
-        String engine,
         String modelType,
-        String tokenizerPath,
-        List<String> specialTokens,
         Integer maxTokens,
-        Integer sequenceLength,
         Integer topk,
         Float temperature,
+        Integer randomSeed,
         ModelLoadCallback callback
     ) {
-        this.modelPath = path;
-        this.tokenizerPath = tokenizerPath;
-        this.modelType = modelType;
-        this.maxTokens = maxTokens;
-        this.sequenceLength = sequenceLength != null ? sequenceLength : maxTokens;
-        this.topk = topk;
-        this.temperature = temperature;
-        try {
-            this.currentEngine = resolveEngine(engine, path, tokenizerPath);
-        } catch (IllegalArgumentException e) {
-            isReady = false;
-            if (callback != null) {
-                callback.onError(e.getMessage());
-            }
-            return;
-        }
-        isReady = false;
+        ModelLoadRequest request = nextLoadRequest(path, modelType, maxTokens, topk, temperature, randomSeed);
+        resetModelState();
 
         android.util.Log.d(
             "LLM",
-            "setModel called with engine: " +
-                currentEngine +
-                ", path: " +
-                path +
-                ", tokenizerPath: " +
-                tokenizerPath +
+            "setModel called with path: " +
+                request.modelPath +
+                ", modelType: " +
+                request.modelType +
                 ", maxTokens: " +
-                maxTokens +
-                ", sequenceLength: " +
-                this.sequenceLength +
+                request.maxTokens +
                 ", topk: " +
-                topk +
+                request.topk +
                 ", temperature: " +
-                temperature +
-                ", specialTokens: " +
-                (specialTokens == null ? 0 : specialTokens.size())
+                request.temperature
         );
-        initializeModel(callback);
+        initializeModel(request, callback);
     }
 
-    private Engine resolveEngine(String engine, String path, String tokenizerPath) {
-        String normalizedEngine = engine == null ? "auto" : engine.toLowerCase(Locale.US);
-        if ("executorch".equals(normalizedEngine)) {
-            return Engine.EXECUTORCH;
-        }
-        if ("mediapipe".equals(normalizedEngine)) {
-            return Engine.MEDIAPIPE;
-        }
-        if ("apple".equals(normalizedEngine)) {
-            throw new IllegalArgumentException("Apple Intelligence is only available on iOS");
-        }
-        if (tokenizerPath != null && !tokenizerPath.isEmpty()) {
-            return Engine.EXECUTORCH;
-        }
-        if (path != null && path.toLowerCase(Locale.US).endsWith(".pte")) {
-            return Engine.EXECUTORCH;
-        }
-        return Engine.MEDIAPIPE;
+    private synchronized ModelLoadRequest nextLoadRequest(
+        String path,
+        String modelType,
+        Integer maxTokens,
+        Integer topk,
+        Float temperature,
+        Integer randomSeed
+    ) {
+        currentLoadId += 1;
+        return new ModelLoadRequest(path, modelType, maxTokens, topk, temperature, randomSeed, currentLoadId);
     }
 
-    private void releaseCurrentModel() {
-        if (mediaPipeInference != null) {
-            try {
-                mediaPipeInference.close();
-            } catch (Exception e) {
-                android.util.Log.w("LLM", "Failed to close MediaPipe inference: " + e.getMessage());
-            } finally {
-                mediaPipeInference = null;
-            }
-        }
-
-        if (executorchModule != null) {
-            try {
-                executorchModule.stop();
-            } catch (Exception e) {
-                android.util.Log.w("LLM", "Failed to stop ExecuTorch module: " + e.getMessage());
-            } finally {
-                executorchModule = null;
-            }
-        }
+    private synchronized boolean isCurrentLoad(long loadId) {
+        return loadId == currentLoadId;
     }
 
-    private void initializeModel(ModelLoadCallback callback) {
-        if (modelPath == null) {
+    private void initializeModel(ModelLoadRequest request, ModelLoadCallback callback) {
+        if (request.modelPath == null) {
             if (callback != null) {
                 callback.onError("Model path not set");
             }
@@ -156,277 +131,530 @@ public class LLM {
 
         executor.execute(() -> {
             try {
-                releaseCurrentModel();
-                if (currentEngine == Engine.EXECUTORCH) {
-                    initializeExecutorchModel();
+                if (!isCurrentLoad(request.loadId)) {
+                    notifySupersededLoad(callback);
+                    return;
+                }
+
+                BackendType backend = resolveBackend(request.modelPath, request.modelType);
+                String actualPath = backend == BackendType.GEMINI_NANO ? request.modelPath : resolveModelPath(request.modelPath, backend);
+
+                android.util.Log.d("LLM", "Resolved model path: " + actualPath + " using backend: " + backend);
+
+                if (backend == BackendType.LITERT) {
+                    Engine loadedEngine = initializeLiteRtModel(actualPath, request);
+                    if (!isCurrentLoad(request.loadId)) {
+                        closeEngine(loadedEngine);
+                        notifySupersededLoad(callback);
+                        return;
+                    }
+                    applyLoadedModel(request, backend, loadedEngine, null, null);
                 } else {
-                    initializeMediaPipeModel();
+                    if (backend == BackendType.GEMINI_NANO) {
+                        initializeGeminiNanoModel(request, callback);
+                        return;
+                    }
+
+                    LlmInference loadedInference = initializeMediaPipeModel(actualPath, request);
+                    if (!isCurrentLoad(request.loadId)) {
+                        closeLlmInference(loadedInference);
+                        notifySupersededLoad(callback);
+                        return;
+                    }
+                    applyLoadedModel(request, backend, null, loadedInference, null);
                 }
-                isReady = true;
 
                 if (callback != null) {
                     callback.onSuccess();
                 }
-            } catch (Exception e) {
-                isReady = false;
-                e.printStackTrace();
+            } catch (Exception exception) {
+                if (!isCurrentLoad(request.loadId)) {
+                    notifySupersededLoad(callback);
+                    return;
+                }
+
+                android.util.Log.e("LLM", "Failed to initialize model", exception);
+                resetModelState();
+
                 if (callback != null) {
-                    callback.onError(e.getMessage());
+                    callback.onError(exception.getMessage());
                 }
             }
         });
     }
 
-    private void initializeMediaPipeModel() throws Exception {
-        String actualPath = resolveModelPath(modelPath, true);
-        android.util.Log.d("LLM", "Final MediaPipe model path: " + actualPath);
+    private void notifySupersededLoad(ModelLoadCallback callback) {
+        if (callback != null) {
+            callback.onError("Model load was superseded by a newer request");
+        }
+    }
 
-        LlmInferenceOptions.Builder optionsBuilder = LlmInferenceOptions.builder()
-            .setModelPath(actualPath)
-            .setMaxTokens(maxTokens)
-            .setMaxTopK(topk);
+    private synchronized void applyLoadedModel(
+        ModelLoadRequest request,
+        BackendType backend,
+        Engine loadedEngine,
+        LlmInference loadedInference,
+        GenerativeModelFutures loadedGeminiNanoModel
+    ) {
+        modelPath = request.modelPath;
+        modelType = request.modelType;
+        maxTokens = request.maxTokens;
+        topk = request.topk;
+        temperature = request.temperature;
+        randomSeed = request.randomSeed;
+        activeBackend = backend;
+        engine = loadedEngine;
+        llmInference = loadedInference;
+        geminiNanoModel = loadedGeminiNanoModel;
+        isReady = true;
+    }
+
+    private Engine initializeLiteRtModel(String actualPath, ModelLoadRequest request) {
+        EngineConfig config = new EngineConfig(
+            actualPath,
+            new Backend.CPU(),
+            new Backend.CPU(),
+            new Backend.CPU(),
+            request.maxTokens,
+            context.getCacheDir().getAbsolutePath()
+        );
+        Engine loadedEngine = new Engine(config);
+        loadedEngine.initialize();
+        return loadedEngine;
+    }
 
-        LlmInferenceOptions options = optionsBuilder.build();
-        mediaPipeInference = LlmInference.createFromOptions(context, options);
+    private LlmInference initializeMediaPipeModel(String actualPath, ModelLoadRequest request) {
+        LlmInferenceOptions options = LlmInferenceOptions.builder()
+            .setModelPath(actualPath)
+            .setMaxTokens(request.maxTokens)
+            .setMaxTopK(request.topk)
+            .build();
+        return LlmInference.createFromOptions(context, options);
     }
 
-    private void initializeExecutorchModel() throws Exception {
-        if (tokenizerPath == null || tokenizerPath.isEmpty()) {
-            throw new RuntimeException("tokenizerPath is required for ExecuTorch models");
+    private void initializeGeminiNanoModel(ModelLoadRequest request, ModelLoadCallback callback) {
+        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.O) {
+            if (callback != null) {
+                callback.onError("Gemini Nano requires Android API 26 or higher");
+            }
+            return;
         }
 
-        String actualModelPath = resolveModelPath(modelPath, false);
-        String actualTokenizerPath = resolveModelPath(tokenizerPath, false);
-        int executorModelType = resolveExecutorchModelType(modelType);
+        GenerativeModelFutures loadedModel = GenerativeModelFutures.from(Generation.INSTANCE.getClient());
+        ListenableFuture<Integer> statusFuture = loadedModel.checkStatus();
+        statusFuture.addListener(
+            () -> {
+                try {
+                    if (!isCurrentLoad(request.loadId)) {
+                        notifySupersededLoad(callback);
+                        return;
+                    }
+
+                    Integer status = statusFuture.get();
+                    if (status == FeatureStatus.AVAILABLE) {
+                        applyLoadedModel(request, BackendType.GEMINI_NANO, null, null, loadedModel);
+                        if (callback != null) {
+                            callback.onSuccess();
+                        }
+                        return;
+                    }
+
+                    resetModelState();
+                    if (callback != null) {
+                        callback.onError("Gemini Nano is " + geminiNanoStatusName(status) + " on this device");
+                    }
+                } catch (CancellationException | ExecutionException exception) {
+                    resetModelState();
+                    if (callback != null) {
+                        callback.onError(exception.getMessage());
+                    }
+                } catch (InterruptedException exception) {
+                    Thread.currentThread().interrupt();
+                    resetModelState();
+                    if (callback != null) {
+                        callback.onError(exception.getMessage());
+                    }
+                }
+            },
+            executor
+        );
+    }
 
-        android.util.Log.d("LLM", "Final ExecuTorch model path: " + actualModelPath);
-        android.util.Log.d("LLM", "Final ExecuTorch tokenizer path: " + actualTokenizerPath);
+    private String geminiNanoStatusName(Integer status) {
+        if (status == null) {
+            return "unknown";
+        }
 
-        executorchModule = new LlmModule(executorModelType, actualModelPath, actualTokenizerPath, temperature);
-        executorchModule.load();
+        switch (status) {
+            case FeatureStatus.UNAVAILABLE:
+                return "unavailable";
+            case FeatureStatus.DOWNLOADABLE:
+                return "downloadable but not downloaded";
+            case FeatureStatus.DOWNLOADING:
+                return "downloading";
+            case FeatureStatus.AVAILABLE:
+                return "available";
+            default:
+                return "unknown";
+        }
     }
 
-    private int resolveExecutorchModelType(String modelType) {
-        if (modelType == null) {
-            return LlmModule.MODEL_TYPE_TEXT;
+    private BackendType resolveBackend(String path, String modelType) {
+        if (isGeminiNanoModel(path, modelType)) {
+            return BackendType.GEMINI_NANO;
         }
-        String normalized = modelType.toLowerCase(Locale.US);
-        if ("text-vision".equals(normalized) || "vision".equals(normalized)) {
-            return LlmModule.MODEL_TYPE_TEXT_VISION;
+
+        return "litertlm".equals(normalizeModelType(path, modelType)) ? BackendType.LITERT : BackendType.MEDIAPIPE;
+    }
+
+    private boolean isGeminiNanoModel(String path, String modelType) {
+        String requestedModel = modelType != null && !modelType.isBlank() ? modelType : path;
+        return "gemininano".equals(normalizeSystemModelName(requestedModel));
+    }
+
+    private String normalizeSystemModelName(String value) {
+        return stripUrlSuffix(value).toLowerCase().replace("-", "").replace("_", "").replace(" ", "");
+    }
+
+    private String normalizeModelType(String path, String modelType) {
+        if (modelType != null && !modelType.isBlank()) {
+            return stripUrlSuffix(modelType).toLowerCase();
         }
-        if ("multimodal".equals(normalized)) {
-            return LlmModule.MODEL_TYPE_MULTIMODAL;
+
+        String normalizedPath = stripUrlSuffix(path);
+        int extensionIndex = normalizedPath.lastIndexOf('.');
+        if (extensionIndex == -1 || extensionIndex == normalizedPath.length() - 1) {
+            return "";
         }
-        return LlmModule.MODEL_TYPE_TEXT;
+
+        return normalizedPath.substring(extensionIndex + 1).toLowerCase();
     }
 
-    private String resolveModelPath(String path, boolean copyMediaPipeCompanion) throws Exception {
-        android.util.Log.d("LLM", "Original model path: " + path);
+    private String stripUrlSuffix(String value) {
+        int cutoff = value.length();
+        int queryIndex = value.indexOf('?');
+        int fragmentIndex = value.indexOf('#');
 
-        if (!path.startsWith("/android_asset/")) {
-            return path;
+        if (queryIndex >= 0 && queryIndex < cutoff) {
+            cutoff = queryIndex;
+        }
+        if (fragmentIndex >= 0 && fragmentIndex < cutoff) {
+            cutoff = fragmentIndex;
         }
 
-        String assetPath = path.substring("/android_asset/".length());
-        android.util.Log.d("LLM", "Asset path: " + assetPath);
+        return value.substring(0, cutoff);
+    }
 
-        try (java.io.InputStream is = context.getAssets().open(assetPath)) {
-            android.util.Log.d("LLM", "Asset exists: " + assetPath);
-        } catch (Exception e) {
-            android.util.Log.e("LLM", "Asset not found: " + assetPath, e);
-            throw new RuntimeException("Asset not found: " + assetPath);
+    private String resolveModelPath(String path, BackendType backendType) throws Exception {
+        if (!path.startsWith("/android_asset/")) {
+            return path;
         }
 
-        java.io.File cacheDir = context.getCacheDir();
-        java.io.File modelFile = new java.io.File(cacheDir, assetPath);
-        java.io.File parent = modelFile.getParentFile();
-        if (parent != null) {
-            parent.mkdirs();
-        }
+        String assetPath = path.substring("/android_asset/".length());
+        File cacheDir = context.getCacheDir();
+        File modelFile = new File(cacheDir, assetPath);
 
         copyAssetToFile(assetPath, modelFile);
-        android.util.Log.d("LLM", "Copied asset to: " + modelFile.getAbsolutePath());
 
-        if (copyMediaPipeCompanion && assetPath.endsWith(".task")) {
-            copyOptionalCompanion(assetPath, cacheDir);
+        if (backendType == BackendType.MEDIAPIPE && assetPath.endsWith(".task")) {
+            String companionPath = assetPath.replace(".task", ".litertlm");
+            try {
+                copyAssetToFile(companionPath, new File(cacheDir, companionPath));
+            } catch (Exception exception) {
+                android.util.Log.d("LLM", "No MediaPipe companion file found for " + assetPath);
+            }
         }
 
         return modelFile.getAbsolutePath();
     }
 
-    private void copyOptionalCompanion(String assetPath, java.io.File cacheDir) {
-        String litertlmPath = assetPath.replace(".task", ".litertlm");
-        try (java.io.InputStream litertlmIs = context.getAssets().open(litertlmPath)) {
-            java.io.File litertlmFile = new java.io.File(cacheDir, litertlmPath);
-            java.io.File parent = litertlmFile.getParentFile();
-            if (parent != null) {
-                parent.mkdirs();
-            }
-            copyAssetToFile(litertlmPath, litertlmFile);
-            android.util.Log.d("LLM", "Also copied companion file: " + litertlmFile.getAbsolutePath());
-        } catch (Exception e) {
-            android.util.Log.d("LLM", "No companion .litertlm file found");
+    private void copyAssetToFile(String assetPath, File destination) throws Exception {
+        File parent = destination.getParentFile();
+        if (parent != null && !parent.exists()) {
+            parent.mkdirs();
         }
-    }
 
-    private void copyAssetToFile(String assetPath, java.io.File destFile) throws Exception {
         try (
-            java.io.InputStream is = context.getAssets().open(assetPath);
-            java.io.FileOutputStream fos = new java.io.FileOutputStream(destFile)
+            InputStream inputStream = context.getAssets().open(assetPath);
+            FileOutputStream outputStream = new FileOutputStream(destination)
         ) {
-            byte[] buffer = new byte[4096];
+            byte[] buffer = new byte[8192];
             int length;
-            while ((length = is.read(buffer)) > 0) {
-                fos.write(buffer, 0, length);
+            while ((length = inputStream.read(buffer)) > 0) {
+                outputStream.write(buffer, 0, length);
             }
         }
     }
 
-    public String createChat() {
+    public synchronized String createChat() {
+        if (!isReady) {
+            throw new IllegalStateException("Model not ready");
+        }
+
         String chatId = UUID.randomUUID().toString();
-        ChatSession session = new ChatSession();
+        ChatSession session;
+
+        if (activeBackend == BackendType.LITERT) {
+            if (engine == null) {
+                throw new IllegalStateException("LiteRT-LM engine not initialized");
+            }
+            session = ChatSession.forLiteRt(createLiteRtConversation());
+        } else if (activeBackend == BackendType.GEMINI_NANO) {
+            if (geminiNanoModel == null) {
+                throw new IllegalStateException("Gemini Nano model not initialized");
+            }
+            session = ChatSession.forGeminiNano();
+        } else {
+            if (llmInference == null) {
+                throw new IllegalStateException("MediaPipe model not initialized");
+            }
+            session = ChatSession.forMediaPipe();
+        }
+
         chatSessions.put(chatId, session);
         return chatId;
     }
 
+    private Conversation createLiteRtConversation() {
+        ConversationConfig defaults = new ConversationConfig();
+        SamplerConfig samplerConfig = new SamplerConfig(topk, DEFAULT_TOP_P, temperature.doubleValue(), randomSeed);
+        ConversationConfig config = defaults.copy(
+            defaults.getSystemInstruction(),
+            defaults.getInitialMessages(),
+            defaults.getTools(),
+            samplerConfig,
+            defaults.getAutomaticToolCalling(),
+            defaults.getChannels()
+        );
+        return engine.createConversation(config);
+    }
+
     public void sendMessage(String chatId, String message, MessageCallback callback) {
-        ChatSession session = chatSessions.get(chatId);
-        if (session == null) {
-            callback.onError("Chat session not found");
+        ChatSession session;
+        LlmInference inference;
+        GenerativeModelFutures nanoModel;
+        Integer sessionTopk;
+        Float sessionTemperature;
+        Integer sessionMaxTokens;
+        Integer sessionRandomSeed;
+        synchronized (this) {
+            session = chatSessions.get(chatId);
+            if (session == null) {
+                throw new IllegalStateException("Chat session not found");
+            }
+
+            if (!isReady) {
+                throw new IllegalStateException("Model not ready");
+            }
+
+            inference = llmInference;
+            nanoModel = geminiNanoModel;
+            sessionTopk = topk;
+            sessionTemperature = temperature;
+            sessionMaxTokens = maxTokens;
+            sessionRandomSeed = randomSeed;
+        }
+
+        if (session.backendType == BackendType.LITERT) {
+            sendMessageWithLiteRt(chatId, message, session, callback);
             return;
         }
 
-        if (!isReady) {
-            callback.onError("Model not ready");
+        if (session.backendType == BackendType.GEMINI_NANO) {
+            if (nanoModel == null) {
+                throw new IllegalStateException("Gemini Nano model not ready");
+            }
+
+            sendMessageWithGeminiNano(
+                chatId,
+                message,
+                session,
+                nanoModel,
+                sessionMaxTokens,
+                sessionTopk,
+                sessionTemperature,
+                sessionRandomSeed,
+                callback
+            );
             return;
         }
 
-        if (currentEngine == Engine.EXECUTORCH) {
-            sendExecutorchMessage(chatId, message, session, callback);
-        } else {
-            sendMediaPipeMessage(chatId, message, session, callback);
+        if (inference == null) {
+            throw new IllegalStateException("MediaPipe model not ready");
         }
-    }
 
-    private void sendExecutorchMessage(String chatId, String message, ChatSession session, MessageCallback callback) {
-        final LlmModule module = executorchModule;
-        final int generationSequenceLength = sequenceLength;
-        final float generationTemperature = temperature;
+        sendMessageWithMediaPipe(chatId, message, session, inference, sessionTopk, sessionTemperature, callback);
+    }
 
-        if (module == null) {
-            callback.onError("ExecuTorch model not ready");
-            return;
+    private void sendMessageWithLiteRt(String chatId, String message, ChatSession session, MessageCallback callback) {
+        if (session.conversation == null) {
+            throw new IllegalStateException("LiteRT-LM conversation not found");
         }
 
-        executor.execute(() -> {
-            try {
-                session.addMessage("user", message);
-                String fullPrompt = session.buildPrompt(message);
-                AtomicBoolean completed = new AtomicBoolean(false);
-                StringBuilder fullResponse = new StringBuilder();
-
-                module.resetContext();
-                int errorCode = module.generate(
-                    fullPrompt,
-                    generationSequenceLength,
-                    new LlmCallback() {
-                        @Override
-                        public void onResult(String result) {
-                            if (result == null || result.isEmpty()) {
-                                return;
-                            }
-                            callback.onTextReceived(chatId, result, true);
-                            fullResponse.append(result);
-                        }
+        session.conversation.sendMessageAsync(
+            message,
+            new com.google.ai.edge.litertlm.MessageCallback() {
+                @Override
+                public void onMessage(Message responseMessage) {
+                    String chunk = extractText(responseMessage);
+                    if (!chunk.isEmpty()) {
+                        callback.onTextReceived(chatId, chunk, true);
+                    }
+                }
 
-                        @Override
-                        public void onStats(String stats) {
-                            if (completed.compareAndSet(false, true)) {
-                                session.addMessage("assistant", fullResponse.toString());
-                                callback.onComplete(chatId);
-                            }
-                        }
-                    },
-                    false,
-                    generationTemperature,
-                    0,
-                    0
-                );
-
-                if (errorCode != 0) {
-                    completed.set(true);
-                    callback.onError("ExecuTorch error " + errorCode);
-                } else if (completed.compareAndSet(false, true)) {
-                    session.addMessage("assistant", fullResponse.toString());
+                @Override
+                public void onDone() {
                     callback.onComplete(chatId);
                 }
-            } catch (Exception e) {
-                callback.onError(e.getMessage());
+
+                @Override
+                public void onError(Throwable throwable) {
+                    callback.onError(throwable.getMessage());
+                }
+            },
+            Collections.emptyMap()
+        );
+    }
+
+    private String extractText(Message responseMessage) {
+        StringBuilder builder = new StringBuilder();
+
+        for (Content content : responseMessage.getContents().getContents()) {
+            if (content instanceof Content.Text) {
+                builder.append(((Content.Text) content).getText());
             }
-        });
+        }
+
+        return builder.toString();
     }
 
-    private void sendMediaPipeMessage(String chatId, String message, ChatSession session, MessageCallback callback) {
-        final LlmInference inference = mediaPipeInference;
-        final int sessionTopk = topk;
-        final float sessionTemperature = temperature;
+    private void sendMessageWithGeminiNano(
+        String chatId,
+        String message,
+        ChatSession session,
+        GenerativeModelFutures nanoModel,
+        Integer sessionMaxTokens,
+        Integer sessionTopk,
+        Float sessionTemperature,
+        Integer sessionRandomSeed,
+        MessageCallback callback
+    ) {
+        synchronized (session) {
+            session.addMessage("user", message);
+        }
 
-        if (inference == null) {
-            callback.onError("MediaPipe model not ready");
-            return;
+        GenerateContentRequest.Builder requestBuilder;
+        synchronized (session) {
+            requestBuilder = new GenerateContentRequest.Builder(new TextPart(session.buildPrompt()));
         }
 
+        requestBuilder.setMaxOutputTokens(Math.min(sessionMaxTokens, 256));
+        requestBuilder.setTopK(sessionTopk);
+        requestBuilder.setTemperature(sessionTemperature);
+        requestBuilder.setSeed(sessionRandomSeed);
+
+        StringBuffer fullResponse = new StringBuffer();
+        AtomicBoolean streamed = new AtomicBoolean(false);
+        ListenableFuture<GenerateContentResponse> responseFuture = nanoModel.generateContent(
+            requestBuilder.build(),
+            new StreamingCallback() {
+                @Override
+                public void onNewText(String additionalText) {
+                    if (additionalText == null || additionalText.isEmpty()) {
+                        return;
+                    }
+
+                    streamed.set(true);
+                    fullResponse.append(additionalText);
+                    callback.onTextReceived(chatId, additionalText, true);
+                }
+            }
+        );
+
+        responseFuture.addListener(
+            () -> {
+                try {
+                    GenerateContentResponse response = responseFuture.get();
+                    if (!streamed.get()) {
+                        String text = extractText(response);
+                        if (!text.isEmpty()) {
+                            fullResponse.append(text);
+                            callback.onTextReceived(chatId, text, true);
+                        }
+                    }
+
+                    synchronized (session) {
+                        session.addMessage("assistant", fullResponse.toString());
+                    }
+                    callback.onComplete(chatId);
+                } catch (CancellationException | ExecutionException exception) {
+                    callback.onError(exception.getMessage());
+                } catch (InterruptedException exception) {
+                    Thread.currentThread().interrupt();
+                    callback.onError(exception.getMessage());
+                }
+            },
+            executor
+        );
+    }
+
+    private String extractText(GenerateContentResponse response) {
+        if (response.getCandidates().isEmpty()) {
+            return "";
+        }
+
+        Candidate candidate = response.getCandidates().get(0);
+        return candidate.getText();
+    }
+
+    private void sendMessageWithMediaPipe(
+        String chatId,
+        String message,
+        ChatSession session,
+        LlmInference inference,
+        Integer sessionTopk,
+        Float sessionTemperature,
+        MessageCallback callback
+    ) {
         executor.execute(() -> {
+            LlmInferenceSession inferenceSession = null;
             try {
-                // Add user message to history
                 session.addMessage("user", message);
 
-                // Build the full prompt with chat history
-                String fullPrompt = session.buildPrompt(message);
+                String fullPrompt = session.buildPrompt();
                 android.util.Log.d("LLM", "Full prompt: " + fullPrompt);
 
-                // Create a session with proper options
-                com.google.mediapipe.tasks.genai.llminference.LlmInferenceSession.LlmInferenceSessionOptions sessionOptions =
-                    com.google.mediapipe.tasks.genai.llminference.LlmInferenceSession.LlmInferenceSessionOptions.builder()
-                        .setTopK(sessionTopk)
-                        .setTemperature(sessionTemperature)
-                        .build();
+                LlmInferenceSession.LlmInferenceSessionOptions sessionOptions = LlmInferenceSession.LlmInferenceSessionOptions.builder()
+                    .setTopK(sessionTopk)
+                    .setTemperature(sessionTemperature)
+                    .build();
 
-                com.google.mediapipe.tasks.genai.llminference.LlmInferenceSession inferenceSession =
-                    com.google.mediapipe.tasks.genai.llminference.LlmInferenceSession.createFromOptions(inference, sessionOptions);
+                inferenceSession = LlmInferenceSession.createFromOptions(inference, sessionOptions);
+                LlmInferenceSession activeSession = inferenceSession;
 
-                // Add the query
                 inferenceSession.addQueryChunk(fullPrompt);
 
-                // Use streaming API
                 final StringBuilder fullResponse = new StringBuilder();
 
                 com.google.mediapipe.tasks.genai.llminference.ProgressListener<String> resultListener =
                     new com.google.mediapipe.tasks.genai.llminference.ProgressListener<String>() {
                         private StringBuilder buffer = new StringBuilder();
                         private boolean hasStarted = false;
-                        private final AtomicBoolean completed = new AtomicBoolean(false);
 
                         @Override
                         public void run(String partialResult, boolean done) {
                             android.util.Log.d("LLM", "Partial result: " + partialResult + ", done: " + done);
 
-                            // Accumulate in buffer
                             buffer.append(partialResult);
 
-                            // Process buffer when we have enough content or when done
                             String content = buffer.toString();
-
-                            // Check if we can process some content
                             StringBuilder toSend = new StringBuilder();
                             int i = 0;
                             while (i < content.length()) {
-                                // Check for escape sequence
                                 if (i < content.length() - 1 && content.charAt(i) == '\\' && content.charAt(i + 1) == 'n') {
                                     toSend.append('\n');
                                     i += 2;
                                 } else if (i == content.length() - 1 && content.charAt(i) == '\\' && !done) {
-                                    // We have a backslash at the end and more chunks coming - wait for next chunk
                                     break;
                                 } else {
                                     toSend.append(content.charAt(i));
@@ -434,14 +662,11 @@ public class LLM {
                                 }
                             }
 
-                            // Update buffer to contain only unprocessed content
                             buffer = new StringBuilder(content.substring(i));
 
-                            // Send processed content if any
                             String chunk = toSend.toString();
 
-                            // Remove leading newline from the very first content
-                            if (!hasStarted && chunk.length() > 0) {
+                            if (!hasStarted && !chunk.isEmpty()) {
                                 chunk = chunk.replaceFirst("^\\n", "");
                                 hasStarted = true;
                             }
@@ -451,37 +676,91 @@ public class LLM {
                                 fullResponse.append(chunk);
                             }
 
-                            if (done && completed.compareAndSet(false, true)) {
-                                if (buffer.length() > 0) {
-                                    String remaining = buffer.toString();
+                            if (done && buffer.length() > 0) {
+                                String remaining = buffer.toString();
+                                if (!remaining.isEmpty()) {
                                     callback.onTextReceived(chatId, remaining, true);
                                     fullResponse.append(remaining);
-                                    buffer = new StringBuilder();
                                 }
+                            }
 
+                            if (done) {
                                 session.addMessage("assistant", fullResponse.toString());
                                 callback.onComplete(chatId);
-
-                                try {
-                                    inferenceSession.close();
-                                } catch (Exception e) {
-                                    android.util.Log.e("LLM", "Failed to close session: " + e.getMessage());
-                                }
+                                closeMediaPipeSession(activeSession);
                             }
                         }
                     };
 
                 inferenceSession.generateResponseAsync(resultListener);
-            } catch (Exception e) {
-                callback.onError(e.getMessage());
+            } catch (Exception exception) {
+                closeMediaPipeSession(inferenceSession);
+                callback.onError(exception.getMessage());
             }
         });
     }
 
-    public String getReadiness() {
+    public synchronized String getReadiness() {
         return isReady ? "ready" : "not_ready";
     }
 
+    private synchronized void resetModelState() {
+        clearChatSessions();
+
+        closeEngine(engine);
+        closeLlmInference(llmInference);
+        engine = null;
+        llmInference = null;
+        geminiNanoModel = null;
+        modelPath = null;
+        modelType = null;
+        isReady = false;
+    }
+
+    private void closeEngine(Engine loadedEngine) {
+        if (loadedEngine != null) {
+            try {
+                loadedEngine.close();
+            } catch (Exception exception) {
+                android.util.Log.w("LLM", "Failed to close LiteRT-LM engine", exception);
+            }
+        }
+    }
+
+    private void closeMediaPipeSession(LlmInferenceSession inferenceSession) {
+        if (inferenceSession != null) {
+            try {
+                inferenceSession.close();
+            } catch (Exception exception) {
+                android.util.Log.e("LLM", "Failed to close MediaPipe session", exception);
+            }
+        }
+    }
+
+    private void closeLlmInference(LlmInference inference) {
+        if (inference != null) {
+            try {
+                inference.close();
+            } catch (Exception exception) {
+                android.util.Log.w("LLM", "Failed to close MediaPipe model", exception);
+            }
+        }
+    }
+
+    private void clearChatSessions() {
+        for (ChatSession session : chatSessions.values()) {
+            if (session.conversation != null) {
+                try {
+                    session.conversation.close();
+                } catch (Exception exception) {
+                    android.util.Log.w("LLM", "Failed to close LiteRT-LM conversation", exception);
+                }
+            }
+        }
+
+        chatSessions.clear();
+    }
+
     public interface MessageCallback {
         void onTextReceived(String chatId, String text, boolean isChunk);
         void onComplete(String chatId);
@@ -493,17 +772,31 @@ public class LLM {
         void onError(String error);
     }
 
-    // Inner class to manage chat sessions
     private static class ChatSession {
 
-        private StringBuilder history;
+        private final BackendType backendType;
+        private final Conversation conversation;
+        private final StringBuilder history;
 
-        ChatSession() {
+        private ChatSession(BackendType backendType, Conversation conversation) {
+            this.backendType = backendType;
+            this.conversation = conversation;
             this.history = new StringBuilder();
         }
 
+        static ChatSession forLiteRt(Conversation conversation) {
+            return new ChatSession(BackendType.LITERT, conversation);
+        }
+
+        static ChatSession forMediaPipe() {
+            return new ChatSession(BackendType.MEDIAPIPE, null);
+        }
+
+        static ChatSession forGeminiNano() {
+            return new ChatSession(BackendType.GEMINI_NANO, null);
+        }
+
         void addMessage(String role, String content) {
-            // Store messages in a cleaner format for history
             if (history.length() > 0) {
                 history.append("\n");
             }
@@ -514,8 +807,37 @@ public class LLM {
             }
         }
 
-        String buildPrompt(String newMessage) {
+        String buildPrompt() {
             return history.toString();
         }
     }
+
+    private static class ModelLoadRequest {
+
+        private final String modelPath;
+        private final String modelType;
+        private final Integer maxTokens;
+        private final Integer topk;
+        private final Float temperature;
+        private final Integer randomSeed;
+        private final long loadId;
+
+        private ModelLoadRequest(
+            String modelPath,
+            String modelType,
+            Integer maxTokens,
+            Integer topk,
+            Float temperature,
+            Integer randomSeed,
+            long loadId
+        ) {
+            this.modelPath = modelPath;
+            this.modelType = modelType;
+            this.maxTokens = maxTokens;
+            this.topk = topk;
+            this.temperature = temperature;
+            this.randomSeed = randomSeed;
+            this.loadId = loadId;
+        }
+    }
 }