react-native-litert-lm 0.2.1 → 0.3.0

package/cpp/HybridLiteRTLM.hpp

@@ -14,11 +14,19 @@
 
 #include "../nitrogen/generated/shared/c++/HybridLiteRTLMSpec.hpp"
 
-// LiteRT-LM headers (conditionally included when available via Prefab/CMake)
-#ifdef LITERT_LM_ENABLED
-#include "litert/lm/engine.h"
-#include "litert/lm/conversation.h"
-#include "litert/lm/types.h"
+// LiteRT-LM C API (iOS uses prebuilt framework with C ABI)
+#ifdef __APPLE__
+#include "include/litert_lm_engine.h"
+#endif
+
+// Memory usage headers
+#ifdef __APPLE__
+#include <mach/mach.h>
+#include <mach/mach_host.h>
+#endif
+#ifdef __ANDROID__
+#include <malloc.h>
+#include <fstream>
 #endif
 
 #include <string>
@@ -27,14 +35,15 @@
 #include <memory>
 #include <mutex>
 #include <functional>
+#include <atomic>
 
 namespace margelo::nitro::litertlm {
 
 /**
  * HybridLiteRTLM: React Native bindings for LiteRT-LM.
  * 
- * Wraps LiteRT-LM's Engine and Conversation classes to provide
- * high-level LLM inference with GPU acceleration.
+ * On iOS, wraps the LiteRT-LM C API (engine.h) with prebuilt framework.
+ * On Android, this class is unused — the Kotlin implementation is used instead.
  */
 class HybridLiteRTLM : public HybridLiteRTLMSpec {
 public:
@@ -51,24 +60,26 @@ public:
 public:
   // HybridLiteRTLMSpec interface implementation
   
-  void loadModel(const std::string& modelPath, 
+  std::shared_ptr<Promise<void>> loadModel(const std::string& modelPath, 
                  const std::optional<LLMConfig>& config) override;
   
-  std::string sendMessage(const std::string& message) override;
+  std::shared_ptr<Promise<std::string>> sendMessage(const std::string& message) override;
   
-  std::string sendMessageWithImage(const std::string& message,
+  std::shared_ptr<Promise<std::string>> sendMessageWithImage(const std::string& message,
                                    const std::string& imagePath) override;
 
-  std::future<std::string> downloadModel(const std::string& url, 
+  std::shared_ptr<Promise<std::string>> downloadModel(const std::string& url, 
                                          const std::string& fileName,
                                          const std::optional<std::function<void(double)>>& onProgress) override;
   
-  std::string sendMessageWithAudio(const std::string& message,
+  std::shared_ptr<Promise<void>> deleteModel(const std::string& fileName) override;
+  
+  std::shared_ptr<Promise<std::string>> sendMessageWithAudio(const std::string& message,
                                    const std::string& audioPath) override;
   
   void sendMessageAsync(
     const std::string& message,
-    const std::function<void(std::string, bool)>& onToken
+    const std::function<void(const std::string&, bool)>& onToken
   ) override;
   
   std::vector<Message> getHistory() override;
@@ -79,13 +90,17 @@ public:
   
   GenerationStats getStats() override;
   
+  MemoryUsage getMemoryUsage() override;
+  
   void close() override;
 
 private:
-  // LiteRT-LM resources (conditionally available on Android with Prefab)
-#ifdef LITERT_LM_ENABLED
-  std::unique_ptr<litert::lm::Engine> engine_;
-  std::unique_ptr<litert::lm::Conversation> conversation_;
+  // LiteRT-LM C API resources (iOS only)
+#ifdef __APPLE__
+  LiteRtLmEngine* engine_ = nullptr;
+  LiteRtLmConversation* conversation_ = nullptr;
+  LiteRtLmConversationConfig* conv_config_ = nullptr;
+  LiteRtLmSessionConfig* session_config_ = nullptr;
 #endif
   
   // State
@@ -96,10 +111,11 @@ private:
   // Thread safety
   mutable std::mutex mutex_;
   
-  // Configuration - backends
+  // Configuration - backend
   Backend backend_ = Backend::GPU;
-  Backend visionBackend_ = Backend::GPU;  // Gemma 3n requires GPU for vision
-  Backend audioBackend_ = Backend::CPU;   // Audio typically CPU
+  
+  // System prompt / instruction
+  std::string systemPrompt_;
   
   // Configuration - sampling parameters
   double temperature_ = 0.7;
@@ -114,11 +130,38 @@ private:
     }
   }
   
-  // Helper to format a message for the engine (apply chat template if needed)
-  std::string formatUserPrompt(const std::string& message) const;
-  
   // Helper to create a new conversation from existing engine
   void createNewConversation();
+  
+  // JSON helpers for building C API message payloads
+  static std::string escapeJson(const std::string& input);
+  static std::string buildTextMessageJson(const std::string& text);
+  static std::string buildImageMessageJson(const std::string& text, const std::string& imagePath);
+  static std::string buildAudioMessageJson(const std::string& text, const std::string& audioPath);
+  static std::string extractTextFromResponse(const std::string& jsonResponse);
+  
+  // Internal implementations (called from Promise lambdas)
+  void loadModelInternal(const std::string& modelPath, const std::optional<LLMConfig>& config);
+  std::string sendMessageInternal(const std::string& message);
+  std::string sendMessageWithImageInternal(const std::string& message, const std::string& imagePath);
+  std::string sendMessageWithAudioInternal(const std::string& message, const std::string& audioPath);
+  
+  // Streaming callback context (must be a plain struct for C function pointer)
+  struct StreamContext {
+    std::function<void(const std::string&, bool)> onToken;
+    std::string fullResponse;
+    std::vector<Message>* history;
+    std::mutex* historyMutex;
+    std::string userMessage;
+    GenerationStats* lastStats;
+    std::chrono::steady_clock::time_point startTime;
+    int tokenCount;
+  };
+  
+  // Static C callback for streaming (no captures needed)
+  static void streamCallbackFn(void* callback_data, const char* chunk,
+                                bool is_final, const char* error_msg);
 };
 
 } // namespace margelo::nitro::litertlm
+

package/cpp/IOSDownloadHelper.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <string>
+#include <functional>
+#include <optional>
+
+namespace litert_lm {
+
+/**
+ * Download a file from a URL to the app's Caches/litert_models directory.
+ * Uses NSURLSession for efficient, resumable downloads.
+ * 
+ * @param url HTTPS URL to download from
+ * @param fileName Destination filename
+ * @param onProgress Optional progress callback (0.0 to 1.0)
+ * @return Absolute path to the downloaded file
+ * @throws std::runtime_error on download failure
+ */
+std::string downloadModelFile(
+    const std::string& url,
+    const std::string& fileName,
+    const std::optional<std::function<void(double)>>& onProgress);
+
+} // namespace litert_lm

package/cpp/cpp-adapter.cpp

@@ -2,12 +2,20 @@
 /// cpp-adapter.cpp
 /// JNI Entry Point - Required by Nitrogen to register Kotlin HybridObjects
 ///
+/// Updated for react-native-nitro-modules v0.35+:
+/// Uses facebook::jni::initialize() directly with registerAllNatives().
+///
 
 #include <jni.h>
+#include <fbjni/fbjni.h>
 #include "LiteRTLMOnLoad.hpp"
 
 // JNI_OnLoad is called when the native library is loaded via System.loadLibrary()
-// This is where we initialize the Nitrogen bridge and register all Kotlin HybridObjects
+// This is where we initialize the Nitrogen bridge and register all Kotlin HybridObjects.
+// The new v0.35 API allows registering custom C++ native JNI classes/functions
+// alongside Nitrogen's auto-generated registrations.
 JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void*) {
-    return margelo::nitro::litertlm::initialize(vm);
+    return facebook::jni::initialize(vm, []() {
+        margelo::nitro::litertlm::registerAllNatives();
+    });
 }