@dr33m/react-native-litert-lm 0.5.0

package/src/inferenceRouting.ts

@@ -0,0 +1,80 @@
+import type { MultimodalPart, PartType } from "./specs/LiteRTLM.nitro";
+
+export type TokenCallback = (token: string, done: boolean) => void;
+
+/** Legacy inference entry points on LiteRTLM (createLLM routes these to execute). */
+export const LEGACY_INFERENCE_METHODS = [
+  "sendMessage",
+  "sendMessageWithImage",
+  "sendMessageWithAudio",
+  "sendMultimodalMessage",
+  "sendMessageAsync",
+  "sendMessageWithImageAsync",
+  "sendMessageWithAudioAsync",
+] as const;
+
+export type LegacyInferenceMethod = (typeof LEGACY_INFERENCE_METHODS)[number];
+
+export function textPart(text: string): MultimodalPart {
+  return { type: "text" as PartType, text };
+}
+
+export function imagePart(path: string): MultimodalPart {
+  return { type: "image" as PartType, path };
+}
+
+export function audioPart(path: string): MultimodalPart {
+  return { type: "audio" as PartType, path };
+}
+
+export type InferenceRoute = {
+  parts: MultimodalPart[];
+  onToken?: TokenCallback;
+};
+
+type LegacyRouteBuilder = (args: unknown[]) => InferenceRoute;
+
+const LEGACY_ROUTES: Record<LegacyInferenceMethod, LegacyRouteBuilder> = {
+  sendMessage: (args) => ({ parts: [textPart(args[0] as string)] }),
+  sendMessageWithImage: (args) => ({
+    parts: [textPart(args[0] as string), imagePart(args[1] as string)],
+  }),
+  sendMessageWithAudio: (args) => ({
+    parts: [textPart(args[0] as string), audioPart(args[1] as string)],
+  }),
+  sendMultimodalMessage: (args) => ({ parts: args[0] as MultimodalPart[] }),
+  sendMessageAsync: (args) => ({
+    parts: [textPart(args[0] as string)],
+    onToken: args[1] as TokenCallback,
+  }),
+  sendMessageWithImageAsync: (args) => ({
+    parts: [textPart(args[0] as string), imagePart(args[1] as string)],
+    onToken: args[2] as TokenCallback,
+  }),
+  sendMessageWithAudioAsync: (args) => ({
+    parts: [textPart(args[0] as string), audioPart(args[1] as string)],
+    onToken: args[2] as TokenCallback,
+  }),
+};
+
+const LEGACY_METHOD_SET = new Set<string>(LEGACY_INFERENCE_METHODS);
+
+/**
+ * Maps legacy public API names to unified MultimodalPart payloads.
+ * JS source of truth — native legacy wrappers mirror these shapes for direct hybrid access.
+ */
+export function routeLegacyInference(
+  method: string,
+  args: unknown[],
+): InferenceRoute | null {
+  if (!LEGACY_METHOD_SET.has(method)) {
+    return null;
+  }
+  return LEGACY_ROUTES[method as LegacyInferenceMethod](args);
+}
+
+export function isLegacyInferenceMethod(
+  method: string,
+): method is LegacyInferenceMethod {
+  return LEGACY_METHOD_SET.has(method);
+}

package/src/memoryTracker.ts

@@ -0,0 +1,268 @@
+/**
+ * Memory tracking utilities for LiteRT-LM using real native memory metrics.
+ *
+ * Records real memory usage from OS-level APIs via `getMemoryUsage()`,
+ * and stores snapshots in a native-backed ArrayBuffer allocated via
+ * `NitroModules.createNativeArrayBuffer()` (v0.35+) for zero-copy interop.
+ *
+ * @example
+ * ```typescript
+ * import { createMemoryTracker } from 'react-native-litert-lm';
+ *
+ * const tracker = createMemoryTracker(100);
+ *
+ * // Record a real snapshot (typically called internally after inference)
+ * tracker.record({
+ *   timestamp: Date.now(),
+ *   nativeHeapBytes: usage.nativeHeapBytes,
+ *   residentBytes: usage.residentBytes,
+ *   availableMemoryBytes: usage.availableMemoryBytes,
+ * });
+ *
+ * console.log(`Peak RSS: ${tracker.getPeakMemory()} bytes`);
+ * ```
+ */
+
+import { NitroModules } from "react-native-nitro-modules";
+
+/**
+ * A single memory usage snapshot with real data from OS APIs.
+ */
+export interface MemorySnapshot {
+  /** Unix timestamp in milliseconds */
+  timestamp: number;
+  /** Native heap allocated bytes (Debug.getNativeHeapAllocatedSize on Android, task_info on iOS) */
+  nativeHeapBytes: number;
+  /** Process resident set size (RSS) in bytes */
+  residentBytes: number;
+  /** Available system memory in bytes */
+  availableMemoryBytes: number;
+}
+
+/** Number of Float64 fields per snapshot */
+const FIELDS_PER_SNAPSHOT = 4;
+/** Bytes per Float64 value */
+const BYTES_PER_FIELD = Float64Array.BYTES_PER_ELEMENT; // 8
+
+/**
+ * Memory tracker that stores snapshots in a native-backed ArrayBuffer.
+ *
+ * Uses `NitroModules.createNativeArrayBuffer()` to allocate the backing
+ * buffer in native (C++) memory, ensuring zero-copy interop with native
+ * methods and keeping memory tracking data off the JS heap.
+ */
+export interface MemoryTracker {
+  /**
+   * Record a new memory snapshot.
+   * @param snapshot The memory usage data to record
+   * @returns true if recorded, false if buffer is full
+   */
+  record(snapshot: MemorySnapshot): boolean;
+
+  /**
+   * Get all recorded snapshots as structured objects.
+   */
+  getSnapshots(): MemorySnapshot[];
+
+  /**
+   * Get the number of recorded snapshots.
+   */
+  getSnapshotCount(): number;
+
+  /**
+   * Get the maximum number of snapshots this tracker can hold.
+   */
+  getCapacity(): number;
+
+  /**
+   * Get the peak resident set size across all snapshots.
+   */
+  getPeakMemory(): number;
+
+  /**
+   * Get the latest memory snapshot, or undefined if none recorded.
+   */
+  getLatestSnapshot(): MemorySnapshot | undefined;
+
+  /**
+   * Get the underlying native ArrayBuffer.
+   * This buffer is allocated via `NitroModules.createNativeArrayBuffer()`
+   * and lives in native memory, enabling zero-copy transfer to native methods.
+   */
+  getNativeBuffer(): ArrayBuffer;
+
+  /**
+   * Get the Float64Array view over the native buffer.
+   */
+  getView(): Float64Array;
+
+  /**
+   * Reset the tracker, clearing all recorded snapshots.
+   * The native buffer is preserved (not reallocated).
+   */
+  reset(): void;
+
+  /**
+   * Get a summary of memory usage statistics.
+   */
+  getSummary(): MemoryTrackerSummary;
+}
+
+/**
+ * Summary statistics from the memory tracker.
+ */
+export interface MemoryTrackerSummary {
+  /** Number of snapshots recorded */
+  snapshotCount: number;
+  /** Peak resident set size in bytes */
+  peakResidentBytes: number;
+  /** Average resident set size in bytes */
+  averageResidentBytes: number;
+  /** Latest resident set size in bytes */
+  currentResidentBytes: number;
+  /** Peak native heap allocated in bytes */
+  peakNativeHeapBytes: number;
+  /** Latest native heap allocated in bytes */
+  currentNativeHeapBytes: number;
+  /** RSS delta from first to last snapshot in bytes */
+  residentDeltaBytes: number;
+  /** Size of the native tracking buffer itself in bytes */
+  trackerBufferSizeBytes: number;
+}
+
+/**
+ * Create a new memory tracker backed by a native ArrayBuffer.
+ *
+ * @param maxSnapshots Maximum number of snapshots to store (default: 256)
+ * @returns A MemoryTracker instance
+ */
+export function createMemoryTracker(maxSnapshots: number = 256): MemoryTracker {
+  const bufferSize = maxSnapshots * FIELDS_PER_SNAPSHOT * BYTES_PER_FIELD;
+
+  // Use NitroModules.createNativeArrayBuffer for native-backed allocation.
+  const nativeBuffer = NitroModules.createNativeArrayBuffer(bufferSize);
+  const view = new Float64Array(nativeBuffer);
+
+  let currentIndex = 0;
+
+  return {
+    record(snapshot: MemorySnapshot): boolean {
+      if (currentIndex >= maxSnapshots) {
+        return false;
+      }
+
+      const offset = currentIndex * FIELDS_PER_SNAPSHOT;
+      view[offset] = snapshot.timestamp;
+      view[offset + 1] = snapshot.nativeHeapBytes;
+      view[offset + 2] = snapshot.residentBytes;
+      view[offset + 3] = snapshot.availableMemoryBytes;
+      currentIndex++;
+
+      return true;
+    },
+
+    getSnapshots(): MemorySnapshot[] {
+      const snapshots: MemorySnapshot[] = [];
+      for (let i = 0; i < currentIndex; i++) {
+        const offset = i * FIELDS_PER_SNAPSHOT;
+        snapshots.push({
+          timestamp: view[offset]!,
+          nativeHeapBytes: view[offset + 1]!,
+          residentBytes: view[offset + 2]!,
+          availableMemoryBytes: view[offset + 3]!,
+        });
+      }
+      return snapshots;
+    },
+
+    getSnapshotCount(): number {
+      return currentIndex;
+    },
+
+    getCapacity(): number {
+      return maxSnapshots;
+    },
+
+    getPeakMemory(): number {
+      let peak = 0;
+      for (let i = 0; i < currentIndex; i++) {
+        const rss = view[i * FIELDS_PER_SNAPSHOT + 2]!;
+        if (rss > peak) {
+          peak = rss;
+        }
+      }
+      return peak;
+    },
+
+    getLatestSnapshot(): MemorySnapshot | undefined {
+      if (currentIndex === 0) return undefined;
+      const offset = (currentIndex - 1) * FIELDS_PER_SNAPSHOT;
+      return {
+        timestamp: view[offset]!,
+        nativeHeapBytes: view[offset + 1]!,
+        residentBytes: view[offset + 2]!,
+        availableMemoryBytes: view[offset + 3]!,
+      };
+    },
+
+    getNativeBuffer(): ArrayBuffer {
+      return nativeBuffer;
+    },
+
+    getView(): Float64Array {
+      return view;
+    },
+
+    reset(): void {
+      view.fill(0);
+      currentIndex = 0;
+    },
+
+    getSummary(): MemoryTrackerSummary {
+      let peakRss = 0;
+      let peakHeap = 0;
+      let sumRss = 0;
+      let firstRss = 0;
+      let lastRss = 0;
+      let lastHeap = 0;
+
+      for (let i = 0; i < currentIndex; i++) {
+        const offset = i * FIELDS_PER_SNAPSHOT;
+        const heap = view[offset + 1]!;
+        const rss = view[offset + 2]!;
+
+        if (rss > peakRss) peakRss = rss;
+        if (heap > peakHeap) peakHeap = heap;
+        sumRss += rss;
+        if (i === 0) firstRss = rss;
+        if (i === currentIndex - 1) {
+          lastRss = rss;
+          lastHeap = heap;
+        }
+      }
+
+      return {
+        snapshotCount: currentIndex,
+        peakResidentBytes: peakRss,
+        averageResidentBytes: currentIndex > 0 ? sumRss / currentIndex : 0,
+        currentResidentBytes: lastRss,
+        peakNativeHeapBytes: peakHeap,
+        currentNativeHeapBytes: lastHeap,
+        residentDeltaBytes: lastRss - firstRss,
+        trackerBufferSizeBytes: bufferSize,
+      };
+    },
+  };
+}
+
+/**
+ * Create a native ArrayBuffer for efficient data transfer.
+ *
+ * A convenience wrapper around `NitroModules.createNativeArrayBuffer()`.
+ *
+ * @param size Size in bytes
+ * @returns A native-backed ArrayBuffer
+ */
+export function createNativeBuffer(size: number): ArrayBuffer {
+  return NitroModules.createNativeArrayBuffer(size);
+}

package/src/modelFactory.ts

@@ -0,0 +1,163 @@
+import { NitroModules } from "react-native-nitro-modules";
+import { LiteRTLM, LLMConfig, MultimodalPart, ExecuteResult } from "./specs/LiteRTLM.nitro";
+import { createMemoryTracker, MemoryTracker } from "./memoryTracker";
+import { ModelRegistry } from "./modelRegistry";
+import {
+  isLegacyInferenceMethod,
+  routeLegacyInference,
+  TokenCallback,
+} from "./inferenceRouting";
+
+/**
+ * Extended LiteRT-LM instance with optional memory tracking and
+ * augmented loadModel that accepts a download progress callback.
+ */
+export type LiteRTLMInstance = Omit<LiteRTLM, "loadModel"> & {
+  memoryTracker?: MemoryTracker;
+  loadModel: (
+    pathOrUrl: string,
+    config?: LLMConfig,
+    onDownloadProgress?: (progress: number) => void,
+  ) => Promise<void>;
+};
+
+/**
+ * Creates a new LiteRT-LM inference engine instance.
+ *
+ * Optionally creates a native-backed memory tracker using
+ * `NitroModules.createNativeArrayBuffer()` (v0.35+) for efficient
+ * zero-copy memory usage tracking.
+ *
+ * @param options.enableMemoryTracking Enable automatic memory tracking (default: false)
+ * @param options.maxMemorySnapshots Maximum number of memory snapshots to store (default: 256)
+ */
+export function createLLM(options?: {
+  enableMemoryTracking?: boolean;
+  maxMemorySnapshots?: number;
+}): LiteRTLMInstance {
+  const native = NitroModules.createHybridObject<LiteRTLM>("LiteRTLM");
+
+  const enableTracking = options?.enableMemoryTracking ?? false;
+  const tracker = enableTracking
+    ? createMemoryTracker(options?.maxMemorySnapshots ?? 256)
+    : undefined;
+
+  const recordMemorySnapshot = () => {
+    if (!tracker) return;
+    try {
+      const usage = native.getMemoryUsage();
+      tracker.record({
+        timestamp: Date.now(),
+        nativeHeapBytes: usage.nativeHeapBytes,
+        residentBytes: usage.residentBytes,
+        availableMemoryBytes: usage.availableMemoryBytes,
+      });
+    } catch {
+      // Non-critical
+    }
+  };
+
+  const augmentedLoadModel = async (
+    pathOrUrl: string,
+    config?: LLMConfig,
+    onDownloadProgress?: (progress: number) => void,
+  ) => {
+    console.log(`Resolving model at ${pathOrUrl}...`);
+    const modelPath = await ModelRegistry.resolveModel(pathOrUrl, {
+      onProgress: onDownloadProgress,
+    });
+    console.log(`Model resolved to: ${modelPath}`);
+
+    const result = await native.loadModel(modelPath, config);
+
+    if (tracker) {
+      tracker.reset();
+      recordMemorySnapshot();
+    }
+
+    return result;
+  };
+
+  /** Single JS inference path — always calls native execute(). */
+  const runExecute = (
+    parts: MultimodalPart[],
+    onToken?: TokenCallback,
+  ): Promise<ExecuteResult> => {
+    const processedParts = parts.map((part) => {
+      if (part.path?.startsWith("file://")) {
+        return { ...part, path: part.path.substring(7) };
+      }
+      return part;
+    });
+
+    if (onToken) {
+      const wrapped: TokenCallback = (token, done) => {
+        onToken(token, done);
+        if (done) recordMemorySnapshot();
+      };
+      return native.execute(processedParts, wrapped);
+    }
+    return native.execute(processedParts, undefined).then((result: ExecuteResult) => {
+      recordMemorySnapshot();
+      return result;
+    });
+  };
+
+  return new Proxy(native, {
+    get(target, prop, receiver) {
+      if (typeof prop !== "string") {
+        return Reflect.get(target, prop, receiver);
+      }
+
+      if (prop === "memoryTracker") {
+        return tracker;
+      }
+      if (prop === "loadModel") {
+        return augmentedLoadModel;
+      }
+      if (prop === "execute") {
+        return (parts: MultimodalPart[], onToken?: TokenCallback) =>
+          runExecute(parts, onToken);
+      }
+      if (prop === "stopGeneration") {
+        return () => target.stopGeneration();
+      }
+      if (prop === "resetConversation") {
+        return () => {
+          const result = target.resetConversation();
+          recordMemorySnapshot();
+          return result;
+        };
+      }
+
+      if (prop === "sendToolResponse") {
+        return (responses: any[], onToken?: TokenCallback) => {
+          if (onToken) {
+            const wrapped: TokenCallback = (token, done) => {
+              onToken(token, done);
+              if (done) recordMemorySnapshot();
+            };
+            return native.sendToolResponse(responses, wrapped);
+          }
+          return native.sendToolResponse(responses, undefined).then((result: ExecuteResult) => {
+            recordMemorySnapshot();
+            return result;
+          });
+        };
+      }
+      if (isLegacyInferenceMethod(prop)) {
+        return (...args: unknown[]) => {
+          const route = routeLegacyInference(prop, args)!;
+          const promise = runExecute(route.parts, route.onToken);
+          return prop.endsWith("Async") ? promise.then(() => {}) : promise;
+        };
+      }
+
+      const original = target[prop as keyof LiteRTLM];
+      if (typeof original === "function") {
+        return original.bind(target);
+      }
+      return original;
+    },
+  }) as unknown as LiteRTLMInstance;
+}

package/src/modelPath.ts

@@ -0,0 +1,16 @@
+/**
+ * Canonical path/URL helpers for model files.
+ * Used by ModelRegistry, hooks, and any download/cache logic.
+ */
+
+/** Extract a filename from a URL or file path (query string stripped). */
+export function extractFileName(pathOrUrl: string): string {
+  const urlWithoutQuery = pathOrUrl.split("?")[0];
+  const name = urlWithoutQuery.split("/").pop();
+  return name || "model.bin";
+}
+
+/** Resolve a cache key from a filename, local path, or HTTPS URL. */
+export function resolveModelFileName(pathOrUrl: string): string {
+  return pathOrUrl.includes("/") ? extractFileName(pathOrUrl) : pathOrUrl;
+}

package/src/modelRegistry.ts

@@ -0,0 +1,104 @@
+import { NitroModules } from "react-native-nitro-modules";
+import type { ModelStore, ModelFile } from "./specs/LiteRTLM.nitro";
+import { extractFileName, resolveModelFileName } from "./modelPath";
+
+export type { ModelFile } from "./specs/LiteRTLM.nitro";
+
+export interface ModelDownloadOptions {
+  headers?: Record<string, string>;
+  onProgress?: (progress: number) => void;
+}
+
+const nativeStore = NitroModules.createHybridObject<ModelStore>("ModelStore");
+
+/**
+ * High-performance Model Registry for react-native-litert-lm.
+ *
+ * Coordinates URL sniffing, HTTPS enforcement, atomic downloading with progress,
+ * and cache queries. Serves as the single source of truth for model storage status.
+ */
+export const ModelRegistry = {
+  /**
+   * Check if a model file is cached locally.
+   * Accepts a filename, local path, or HTTPS URL.
+   *
+   * @param pathOrUrl Filename, local path, or download URL
+   * @returns true if cached and has size > 0
+   */
+  isCached(pathOrUrl: string): boolean {
+    return nativeStore.isCached(resolveModelFileName(pathOrUrl));
+  },
+
+  /**
+   * Get the absolute local path of a cached model.
+   * Accepts a filename, local path, or HTTPS URL.
+   *
+   * @param pathOrUrl Filename, local path, or download URL
+   * @returns The absolute local path
+   */
+  getFilePath(pathOrUrl: string): string {
+    return nativeStore.getFilePath(resolveModelFileName(pathOrUrl));
+  },
+
+  /**
+   * List all locally cached model files.
+   *
+   * @returns Array of ModelFile descriptors containing path, size, and mod time
+   */
+  listCachedFiles(): ModelFile[] {
+    return nativeStore.listCachedFiles();
+  },
+
+  /**
+   * Delete a cached model file.
+   * Accepts a filename, local path, or HTTPS URL.
+   *
+   * @param pathOrUrl Filename, local path, or download URL to delete
+   */
+  deleteFile(pathOrUrl: string): void {
+    nativeStore.deleteFile(resolveModelFileName(pathOrUrl));
+  },
+
+  /**
+   * Resolve a model path or URL.
+   *
+   * Sniffs the protocol. If it is an HTTPS URL, downloads it to the native
+   * cache directory (and throttles callbacks) if not already cached. If it is a
+   * local path, returns it directly.
+   *
+   * @param pathOrUrl Local filepath or HTTPS download URL
+   * @param options Custom HTTP headers and progress callback
+   * @returns Promise resolving to the absolute local path of the model
+   */
+  async resolveModel(pathOrUrl: string, options?: ModelDownloadOptions): Promise<string> {
+    const cleanPath = pathOrUrl.startsWith("file://") ? pathOrUrl.substring(7) : pathOrUrl;
+    if (cleanPath.startsWith("http://") || cleanPath.startsWith("https://")) {
+      if (cleanPath.startsWith("http://")) {
+        throw new Error(
+          "Insecure HTTP URLs are not allowed for model downloads. " +
+            "Use HTTPS instead: " +
+            cleanPath.replace("http://", "https://")
+        );
+      }
+
+      const urlWithoutQuery = cleanPath.split("?")[0];
+      const fileName = urlWithoutQuery.split("/").pop();
+      if (!fileName) {
+        throw new Error(`Invalid model URL: ${cleanPath}`);
+      }
+
+      const headersJson = options?.headers ? JSON.stringify(options.headers) : "{}";
+      
+      return nativeStore.downloadFile(
+        cleanPath,
+        fileName,
+        headersJson,
+        (progress) => {
+          options?.onProgress?.(progress);
+        }
+      );
+    }
+
+    return cleanPath;
+  }
+};