react-native-litert-lm 0.3.7 → 0.4.0

package/ios/Tests/HybridLiteRTLMTests.swift

@@ -0,0 +1,67 @@
+import XCTest
+@testable import LiteRTLM
+
+class HybridLiteRTLMTests: XCTestCase {
+    var bridge: HybridLiteRTLM!
+
+    override func setUp() {
+        super.setUp()
+        bridge = HybridLiteRTLM()
+    }
+
+    override func tearDown() {
+        try? bridge.close()
+        bridge = nil
+        super.tearDown()
+    }
+
+    func testPathTraversalRejection() async throws {
+        let traversals = ["../../etc/passwd", "/absolute/path/file", "subdir\\..\\file", "..", "../", "..\\"]
+        for traversal in traversals {
+            do {
+                let promise = try bridge.deleteModel(fileName: traversal)
+                _ = try await promise.await()
+                XCTFail("Should have failed for traversal: \(traversal)")
+            } catch {
+                let nsError = error as NSError
+                XCTAssertEqual(nsError.domain, "LiteRTLM")
+                XCTAssertEqual(nsError.code, 400)
+                XCTAssertTrue(nsError.localizedDescription.contains("path traversal") || nsError.localizedDescription.contains("directory separators"))
+            }
+        }
+    }
+
+    func testNonHTTPSDownloadRejection() async throws {
+        do {
+            let promise = try bridge.downloadModel(url: "http://insecure-domain.com/model.bin", fileName: "model.bin", onProgress: nil)
+            _ = try await promise.await()
+            XCTFail("Should have blocked insecure HTTP downloads")
+        } catch {
+            let nsError = error as NSError
+            XCTAssertEqual(nsError.domain, "LiteRTLM")
+            XCTAssertEqual(nsError.code, 400)
+            XCTAssertTrue(nsError.localizedDescription.contains("HTTPS is required"))
+        }
+    }
+
+    func testMemoryTelemetry() {
+        XCTAssertNoThrow(try bridge.getMemoryUsage())
+        if let mem = try? bridge.getMemoryUsage() {
+            XCTAssertGreaterThanOrEqual(mem.nativeHeapBytes, 0.0)
+            XCTAssertGreaterThanOrEqual(mem.residentBytes, 0.0)
+            XCTAssertGreaterThanOrEqual(mem.availableMemoryBytes, 0.0)
+        }
+    }
+
+    func testInitialStats() {
+        XCTAssertNoThrow(try bridge.getStats())
+        if let stats = try? bridge.getStats() {
+            XCTAssertEqual(stats.promptTokens, 0.0)
+            XCTAssertEqual(stats.completionTokens, 0.0)
+            XCTAssertEqual(stats.totalTokens, 0.0)
+            XCTAssertEqual(stats.timeToFirstToken, 0.0)
+            XCTAssertEqual(stats.totalTime, 0.0)
+            XCTAssertEqual(stats.tokensPerSecond, 0.0)
+        }
+    }
+}

package/lib/__mocks__/react-native-nitro-modules.d.ts

@@ -0,0 +1,61 @@
+export declare const mockLiteRTLM: {
+    isReady: jest.Mock<boolean, [], any>;
+    loadModel: jest.Mock<any, any, any>;
+    sendMessage: jest.Mock<any, any, any>;
+    sendMessageWithImage: jest.Mock<any, any, any>;
+    downloadModel: jest.Mock<Promise<string>, [url: any, fileName: any, onProgress: any], any>;
+    deleteModel: jest.Mock<any, any, any>;
+    sendMessageWithAudio: jest.Mock<any, any, any>;
+    sendMultimodalMessage: jest.Mock<any, any, any>;
+    sendMessageAsync: jest.Mock<Promise<void>, [msg: any, onToken: any], any>;
+    getHistory: jest.Mock<never[], [], any>;
+    resetConversation: jest.Mock<any, any, any>;
+    getStats: jest.Mock<{
+        promptTokens: number;
+        completionTokens: number;
+        totalTokens: number;
+        timeToFirstToken: number;
+        totalTime: number;
+        tokensPerSecond: number;
+    }, [], any>;
+    countTokens: jest.Mock<number, [], any>;
+    getMemoryUsage: jest.Mock<{
+        nativeHeapBytes: number;
+        residentBytes: number;
+        availableMemoryBytes: number;
+        isLowMemory: boolean;
+    }, [], any>;
+    close: jest.Mock<any, any, any>;
+};
+export declare const NitroModules: {
+    createHybridObject: jest.Mock<{
+        isReady: jest.Mock<boolean, [], any>;
+        loadModel: jest.Mock<any, any, any>;
+        sendMessage: jest.Mock<any, any, any>;
+        sendMessageWithImage: jest.Mock<any, any, any>;
+        downloadModel: jest.Mock<Promise<string>, [url: any, fileName: any, onProgress: any], any>;
+        deleteModel: jest.Mock<any, any, any>;
+        sendMessageWithAudio: jest.Mock<any, any, any>;
+        sendMultimodalMessage: jest.Mock<any, any, any>;
+        sendMessageAsync: jest.Mock<Promise<void>, [msg: any, onToken: any], any>;
+        getHistory: jest.Mock<never[], [], any>;
+        resetConversation: jest.Mock<any, any, any>;
+        getStats: jest.Mock<{
+            promptTokens: number;
+            completionTokens: number;
+            totalTokens: number;
+            timeToFirstToken: number;
+            totalTime: number;
+            tokensPerSecond: number;
+        }, [], any>;
+        countTokens: jest.Mock<number, [], any>;
+        getMemoryUsage: jest.Mock<{
+            nativeHeapBytes: number;
+            residentBytes: number;
+            availableMemoryBytes: number;
+            isLowMemory: boolean;
+        }, [], any>;
+        close: jest.Mock<any, any, any>;
+    }, [name: string], any>;
+    createNativeArrayBuffer: jest.Mock<ArrayBuffer, [size: number], any>;
+};

package/lib/__mocks__/react-native-nitro-modules.js

@@ -0,0 +1,50 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.NitroModules = exports.mockLiteRTLM = void 0;
+exports.mockLiteRTLM = {
+    isReady: jest.fn(() => false),
+    loadModel: jest.fn().mockResolvedValue(undefined),
+    sendMessage: jest.fn().mockResolvedValue("Mock response"),
+    sendMessageWithImage: jest.fn().mockResolvedValue("Mock vision response"),
+    downloadModel: jest.fn(async (url, fileName, onProgress) => {
+        onProgress?.(1.0);
+        return "/mock/path/model.litertlm";
+    }),
+    deleteModel: jest.fn().mockResolvedValue(undefined),
+    sendMessageWithAudio: jest.fn().mockResolvedValue("Mock audio response"),
+    sendMultimodalMessage: jest.fn().mockResolvedValue("Mock multimodal response"),
+    sendMessageAsync: jest.fn((msg, onToken) => {
+        onToken("Mock ", false);
+        onToken("token", true);
+        return Promise.resolve();
+    }),
+    getHistory: jest.fn(() => []),
+    resetConversation: jest.fn(),
+    getStats: jest.fn(() => ({
+        promptTokens: 10,
+        completionTokens: 20,
+        totalTokens: 30,
+        timeToFirstToken: 5,
+        totalTime: 50,
+        tokensPerSecond: 400,
+    })),
+    countTokens: jest.fn(() => -1),
+    getMemoryUsage: jest.fn(() => ({
+        nativeHeapBytes: 1000000,
+        residentBytes: 2000000,
+        availableMemoryBytes: 4000000,
+        isLowMemory: false,
+    })),
+    close: jest.fn(),
+};
+exports.NitroModules = {
+    createHybridObject: jest.fn((name) => {
+        if (name === "LiteRTLM") {
+            return exports.mockLiteRTLM;
+        }
+        throw new Error(`Mock not implemented for hybrid object: ${name}`);
+    }),
+    createNativeArrayBuffer: jest.fn((size) => {
+        return new ArrayBuffer(size);
+    }),
+};

package/lib/__tests__/hooks.test.d.ts

@@ -0,0 +1 @@
+export {};

package/lib/__tests__/hooks.test.js

@@ -0,0 +1,124 @@
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+// Configure React act environment
+global.IS_REACT_ACT_ENVIRONMENT = true;
+const hooks_1 = require("../hooks");
+const react_native_nitro_modules_1 = require("../__mocks__/react-native-nitro-modules");
+const react_test_renderer_1 = __importDefault(require("react-test-renderer"));
+const react_1 = __importDefault(require("react"));
+// Mock react-native
+jest.mock('react-native', () => ({
+    Platform: {
+        OS: 'ios',
+        select: jest.fn((dict) => dict.ios),
+    },
+}));
+// Helper to render and test hooks using react-test-renderer
+function renderHook(callback, initialProps) {
+    let result = { current: null };
+    const TestComponent = ({ props }) => {
+        result.current = callback(props);
+        return null;
+    };
+    let renderer;
+    react_test_renderer_1.default.act(() => {
+        renderer = react_test_renderer_1.default.create(react_1.default.createElement(TestComponent, { props: initialProps }));
+    });
+    const rerender = (newProps) => {
+        react_test_renderer_1.default.act(() => {
+            renderer.update(react_1.default.createElement(TestComponent, { props: newProps }));
+        });
+    };
+    const unmount = () => {
+        react_test_renderer_1.default.act(() => {
+            renderer.unmount();
+        });
+    };
+    return { result, rerender, unmount };
+}
+describe('useModel React Hook Unit Tests', () => {
+    beforeEach(() => {
+        jest.clearAllMocks();
+    });
+    it('should initialize with correct default state and call loadModel automatically when autoLoad is true', async () => {
+        let hookResult;
+        await react_test_renderer_1.default.act(async () => {
+            hookResult = renderHook(() => (0, hooks_1.useModel)('https://example.com/model.litertlm', { autoLoad: true }));
+        });
+        expect(hookResult.result.current.isReady).toBe(true);
+        expect(hookResult.result.current.isGenerating).toBe(false);
+        expect(hookResult.result.current.downloadProgress).toBe(1); // loadModel completed
+        expect(hookResult.result.current.error).toBeNull();
+        expect(react_native_nitro_modules_1.mockLiteRTLM.loadModel).toHaveBeenCalled();
+    });
+    it('should not call loadModel automatically when autoLoad is false', async () => {
+        let hookResult;
+        await react_test_renderer_1.default.act(async () => {
+            hookResult = renderHook(() => (0, hooks_1.useModel)('https://example.com/model.litertlm', { autoLoad: false }));
+        });
+        expect(hookResult.result.current.isReady).toBe(false);
+        expect(react_native_nitro_modules_1.mockLiteRTLM.loadModel).not.toHaveBeenCalled();
+        // Call load manually
+        await react_test_renderer_1.default.act(async () => {
+            await hookResult.result.current.load();
+        });
+        expect(hookResult.result.current.isReady).toBe(true);
+        expect(react_native_nitro_modules_1.mockLiteRTLM.loadModel).toHaveBeenCalled();
+    });
+    it('should handle model load failure gracefully', async () => {
+        react_native_nitro_modules_1.mockLiteRTLM.loadModel.mockRejectedValueOnce(new Error("Model load failed"));
+        let hookResult;
+        await react_test_renderer_1.default.act(async () => {
+            hookResult = renderHook(() => (0, hooks_1.useModel)('https://example.com/model.litertlm', { autoLoad: true }));
+        });
+        expect(hookResult.result.current.isReady).toBe(false);
+        expect(hookResult.result.current.error).toBe("Model load failed");
+    });
+    it('should generate text successfully and trigger memory summary update', async () => {
+        let hookResult;
+        await react_test_renderer_1.default.act(async () => {
+            hookResult = renderHook(() => (0, hooks_1.useModel)('https://example.com/model.litertlm', {
+                autoLoad: true,
+                enableMemoryTracking: true
+            }));
+        });
+        let response = "";
+        await react_test_renderer_1.default.act(async () => {
+            response = await hookResult.result.current.generate("Test prompt");
+        });
+        expect(response).toBe("Mock token");
+        expect(react_native_nitro_modules_1.mockLiteRTLM.sendMessageAsync).toHaveBeenCalled();
+        expect(hookResult.result.current.memorySummary).toBeDefined();
+    });
+    it('should reset conversation correctly', async () => {
+        let hookResult;
+        await react_test_renderer_1.default.act(async () => {
+            hookResult = renderHook(() => (0, hooks_1.useModel)('https://example.com/model.litertlm', { autoLoad: true }));
+        });
+        hookResult.result.current.reset();
+        expect(react_native_nitro_modules_1.mockLiteRTLM.resetConversation).toHaveBeenCalled();
+    });
+    it('should delete model file correctly', async () => {
+        let hookResult;
+        await react_test_renderer_1.default.act(async () => {
+            hookResult = renderHook(() => (0, hooks_1.useModel)('https://example.com/model.litertlm', { autoLoad: true }));
+        });
+        await react_test_renderer_1.default.act(async () => {
+            await hookResult.result.current.deleteModel();
+        });
+        expect(react_native_nitro_modules_1.mockLiteRTLM.deleteModel).toHaveBeenCalledWith('model.litertlm');
+        expect(hookResult.result.current.isReady).toBe(false);
+        expect(hookResult.result.current.downloadProgress).toBe(0);
+    });
+    it('should call close on unmount', async () => {
+        let hookResult;
+        await react_test_renderer_1.default.act(async () => {
+            hookResult = renderHook(() => (0, hooks_1.useModel)('https://example.com/model.litertlm', { autoLoad: false }));
+        });
+        hookResult.unmount();
+        expect(react_native_nitro_modules_1.mockLiteRTLM.close).toHaveBeenCalled();
+    });
+});

package/lib/__tests__/memoryTracker.test.d.ts

@@ -0,0 +1 @@
+export {};

package/lib/__tests__/memoryTracker.test.js

@@ -0,0 +1,74 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+const memoryTracker_1 = require("../memoryTracker");
+const react_native_nitro_modules_1 = require("react-native-nitro-modules");
+describe('MemoryTracker Unit Tests', () => {
+    beforeEach(() => {
+        jest.clearAllMocks();
+    });
+    it('should allocate correct native-backed ArrayBuffer size on initialization', () => {
+        const tracker = (0, memoryTracker_1.createMemoryTracker)(10);
+        expect(react_native_nitro_modules_1.NitroModules.createNativeArrayBuffer).toHaveBeenCalledWith(10 * 4 * 8); // 10 snapshots * 4 fields * 8 bytes/Float64
+        expect(tracker.getCapacity()).toBe(10);
+        expect(tracker.getSnapshotCount()).toBe(0);
+    });
+    it('should record snapshots correctly and retrieve them', () => {
+        const tracker = (0, memoryTracker_1.createMemoryTracker)(5);
+        const snapshot1 = {
+            timestamp: 1000,
+            nativeHeapBytes: 100,
+            residentBytes: 200,
+            availableMemoryBytes: 500,
+        };
+        expect(tracker.record(snapshot1)).toBe(true);
+        expect(tracker.getSnapshotCount()).toBe(1);
+        expect(tracker.getLatestSnapshot()).toEqual(snapshot1);
+        const snapshots = tracker.getSnapshots();
+        expect(snapshots).toHaveLength(1);
+        expect(snapshots[0]).toEqual(snapshot1);
+    });
+    it('should reject new snapshots and return false when capacity is reached', () => {
+        const tracker = (0, memoryTracker_1.createMemoryTracker)(2);
+        expect(tracker.record({ timestamp: 1, nativeHeapBytes: 10, residentBytes: 20, availableMemoryBytes: 50 })).toBe(true);
+        expect(tracker.record({ timestamp: 2, nativeHeapBytes: 20, residentBytes: 30, availableMemoryBytes: 40 })).toBe(true);
+        expect(tracker.record({ timestamp: 3, nativeHeapBytes: 30, residentBytes: 40, availableMemoryBytes: 30 })).toBe(false);
+        expect(tracker.getSnapshotCount()).toBe(2);
+    });
+    it('should calculate correct peak resident memory size', () => {
+        const tracker = (0, memoryTracker_1.createMemoryTracker)(5);
+        tracker.record({ timestamp: 1, nativeHeapBytes: 100, residentBytes: 150, availableMemoryBytes: 1000 });
+        tracker.record({ timestamp: 2, nativeHeapBytes: 120, residentBytes: 300, availableMemoryBytes: 1000 });
+        tracker.record({ timestamp: 3, nativeHeapBytes: 110, residentBytes: 200, availableMemoryBytes: 1000 });
+        expect(tracker.getPeakMemory()).toBe(300);
+    });
+    it('should calculate accurate memory summary statistics', () => {
+        const tracker = (0, memoryTracker_1.createMemoryTracker)(5);
+        tracker.record({ timestamp: 1, nativeHeapBytes: 50, residentBytes: 100, availableMemoryBytes: 1000 });
+        tracker.record({ timestamp: 2, nativeHeapBytes: 150, residentBytes: 300, availableMemoryBytes: 800 });
+        tracker.record({ timestamp: 3, nativeHeapBytes: 100, residentBytes: 200, availableMemoryBytes: 900 });
+        const summary = tracker.getSummary();
+        expect(summary.snapshotCount).toBe(3);
+        expect(summary.peakResidentBytes).toBe(300);
+        expect(summary.averageResidentBytes).toBe(200); // (100 + 300 + 200) / 3
+        expect(summary.currentResidentBytes).toBe(200);
+        expect(summary.peakNativeHeapBytes).toBe(150);
+        expect(summary.currentNativeHeapBytes).toBe(100);
+        expect(summary.residentDeltaBytes).toBe(100); // currentRss(200) - firstRss(100)
+        expect(summary.trackerBufferSizeBytes).toBe(5 * 4 * 8);
+    });
+    it('should preserve buffer but reset internal state when reset() is called', () => {
+        const tracker = (0, memoryTracker_1.createMemoryTracker)(5);
+        tracker.record({ timestamp: 1, nativeHeapBytes: 50, residentBytes: 100, availableMemoryBytes: 1000 });
+        expect(tracker.getSnapshotCount()).toBe(1);
+        tracker.reset();
+        expect(tracker.getSnapshotCount()).toBe(0);
+        expect(tracker.getLatestSnapshot()).toBeUndefined();
+        expect(tracker.getSnapshots()).toEqual([]);
+    });
+    it('should allow standalone native ArrayBuffer allocation via createNativeBuffer', () => {
+        const size = 128;
+        const buffer = (0, memoryTracker_1.createNativeBuffer)(size);
+        expect(react_native_nitro_modules_1.NitroModules.createNativeArrayBuffer).toHaveBeenCalledWith(size);
+        expect(buffer.byteLength).toBe(size);
+    });
+});

package/lib/__tests__/modelFactory.test.d.ts

@@ -0,0 +1 @@
+export {};

package/lib/__tests__/modelFactory.test.js

@@ -0,0 +1,52 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+const modelFactory_1 = require("../modelFactory");
+const react_native_nitro_modules_1 = require("../__mocks__/react-native-nitro-modules");
+describe('modelFactory Security & Proxy Unit Tests', () => {
+    let llm;
+    beforeEach(() => {
+        jest.clearAllMocks();
+        llm = (0, modelFactory_1.createLLM)({ enableMemoryTracking: true });
+    });
+    it('should block insecure HTTP downloads', async () => {
+        await expect(llm.loadModel('http://example.com/model.litertlm'))
+            .rejects.toThrow('Insecure HTTP URLs are not allowed for model downloads');
+    });
+    it('should allow secure HTTPS downloads and strip query parameters', async () => {
+        await llm.loadModel('https://example.com/model.litertlm?token=123');
+        expect(react_native_nitro_modules_1.mockLiteRTLM.downloadModel).toHaveBeenCalledWith('https://example.com/model.litertlm?token=123', 'model.litertlm', expect.any(Function));
+        expect(react_native_nitro_modules_1.mockLiteRTLM.loadModel).toHaveBeenCalledWith('/mock/path/model.litertlm', undefined);
+    });
+    it('should throw an error for invalid model URL', async () => {
+        await expect(llm.loadModel('https://example.com/'))
+            .rejects.toThrow('Invalid model URL: https://example.com/');
+    });
+    it('should successfully proxy sendMessage and record memory metrics', async () => {
+        const response = await llm.sendMessage("Test prompt");
+        expect(response).toBe("Mock response");
+        expect(react_native_nitro_modules_1.mockLiteRTLM.sendMessage).toHaveBeenCalledWith("Test prompt");
+        expect(react_native_nitro_modules_1.mockLiteRTLM.getMemoryUsage).toHaveBeenCalled();
+        expect(llm.memoryTracker?.getSnapshotCount()).toBe(1); // sendMessage records one
+    });
+    it('should successfully proxy resetConversation and record memory metrics', async () => {
+        await llm.resetConversation();
+        expect(react_native_nitro_modules_1.mockLiteRTLM.resetConversation).toHaveBeenCalled();
+        expect(react_native_nitro_modules_1.mockLiteRTLM.getMemoryUsage).toHaveBeenCalled();
+    });
+    it('should successfully proxy sendMessageAsync and record memory metrics when done', async () => {
+        const onToken = jest.fn();
+        await llm.sendMessageAsync("Async prompt", onToken);
+        expect(onToken).toHaveBeenCalledWith("Mock ", false);
+        expect(onToken).toHaveBeenCalledWith("token", true);
+        expect(react_native_nitro_modules_1.mockLiteRTLM.sendMessageAsync).toHaveBeenCalled();
+        expect(react_native_nitro_modules_1.mockLiteRTLM.getMemoryUsage).toHaveBeenCalled();
+    });
+    it('should successfully access memoryTracker and getSnapshots when memory tracking is enabled', () => {
+        expect(llm.memoryTracker).toBeDefined();
+        expect(llm.memoryTracker?.getCapacity()).toBe(256);
+    });
+    it('should not initialize memoryTracker when enableMemoryTracking option is false', () => {
+        const untrackedLLM = (0, modelFactory_1.createLLM)({ enableMemoryTracking: false });
+        expect(untrackedLLM.memoryTracker).toBeUndefined();
+    });
+});

package/lib/hooks.js

@@ -105,7 +105,7 @@ function useModel(pathOrUrl, config) {
                             refreshMemorySummary();
                             resolve(fullResponse);
                         }
-                    });
+                    }).catch(reject);
                 }
                 catch (e) {
                     reject(e);

package/lib/index.d.ts

@@ -1,7 +1,5 @@
 import type { Backend } from "./specs/LiteRTLM.nitro";
 export type { LiteRTLM, LLMConfig, Message, Backend, Role, GenerationStats, MemoryUsage, } from "./specs/LiteRTLM.nitro";
-export type { ChatMessage } from "./templates";
-export { applyGemmaTemplate, applyPhiTemplate, applyLlamaTemplate, } from "./templates";
 export type { MemorySnapshot, MemoryTracker, MemoryTrackerSummary, } from "./memoryTracker";
 export { createMemoryTracker, createNativeBuffer } from "./memoryTracker";
 export type { LiteRTLMInstance } from "./modelFactory";

package/lib/index.js

@@ -14,15 +14,11 @@ var __exportStar = (this && this.__exportStar) || function(m, exports) {
     for (var p in m) if (p !== "default" && !Object.prototype.hasOwnProperty.call(exports, p)) __createBinding(exports, m, p);
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.GEMMA_4_E4B_IT = exports.GEMMA_4_E2B_IT = exports.GEMMA_3N_E2B_IT_INT4 = exports.Models = exports.createLLM = exports.createNativeBuffer = exports.createMemoryTracker = exports.applyLlamaTemplate = exports.applyPhiTemplate = exports.applyGemmaTemplate = void 0;
+exports.GEMMA_4_E4B_IT = exports.GEMMA_4_E2B_IT = exports.GEMMA_3N_E2B_IT_INT4 = exports.Models = exports.createLLM = exports.createNativeBuffer = exports.createMemoryTracker = void 0;
 exports.getRecommendedBackend = getRecommendedBackend;
 exports.checkBackendSupport = checkBackendSupport;
 exports.checkMultimodalSupport = checkMultimodalSupport;
 const react_native_1 = require("react-native");
-var templates_1 = require("./templates");
-Object.defineProperty(exports, "applyGemmaTemplate", { enumerable: true, get: function () { return templates_1.applyGemmaTemplate; } });
-Object.defineProperty(exports, "applyPhiTemplate", { enumerable: true, get: function () { return templates_1.applyPhiTemplate; } });
-Object.defineProperty(exports, "applyLlamaTemplate", { enumerable: true, get: function () { return templates_1.applyLlamaTemplate; } });
 var memoryTracker_1 = require("./memoryTracker");
 Object.defineProperty(exports, "createMemoryTracker", { enumerable: true, get: function () { return memoryTracker_1.createMemoryTracker; } });
 Object.defineProperty(exports, "createNativeBuffer", { enumerable: true, get: function () { return memoryTracker_1.createNativeBuffer; } });

package/lib/modelFactory.js

@@ -38,72 +38,71 @@ function createLLM(options) {
             // Ignore errors during memory tracking - it's non-critical
         }
     };
-    return {
-        ...native,
-        memoryTracker: tracker,
-        loadModel: async (pathOrUrl, config, onDownloadProgress) => {
-            let modelPath = pathOrUrl;
-            // Check if it's a URL — enforce HTTPS for model downloads
-            if (pathOrUrl.startsWith("http://") || pathOrUrl.startsWith("https://")) {
-                if (pathOrUrl.startsWith("http://")) {
-                    throw new Error("Insecure HTTP URLs are not allowed for model downloads. " +
-                        "Use HTTPS instead: " +
-                        pathOrUrl.replace("http://", "https://"));
-                }
-                // Extract filename from URL
-                const fileName = pathOrUrl.split("/").pop();
-                if (!fileName) {
-                    throw new Error(`Invalid model URL: ${pathOrUrl}`);
-                }
-                console.log(`Checking model at ${pathOrUrl}...`);
-                modelPath = await native.downloadModel(pathOrUrl, fileName, (progress) => {
-                    onDownloadProgress?.(progress);
-                });
-                console.log(`Model downloaded to: ${modelPath}`);
+    const augmentedLoadModel = async (pathOrUrl, config, onDownloadProgress) => {
+        let modelPath = pathOrUrl;
+        // Check if it's a URL — enforce HTTPS for model downloads
+        if (pathOrUrl.startsWith("http://") || pathOrUrl.startsWith("https://")) {
+            if (pathOrUrl.startsWith("http://")) {
+                throw new Error("Insecure HTTP URLs are not allowed for model downloads. " +
+                    "Use HTTPS instead: " +
+                    pathOrUrl.replace("http://", "https://"));
             }
-            const result = await native.loadModel(modelPath, config);
-            // Record initial memory snapshot after model load
-            if (tracker) {
-                tracker.reset();
-                recordMemorySnapshot();
+            // Extract filename from URL, stripping query parameters
+            const urlWithoutQuery = pathOrUrl.split("?")[0];
+            const fileName = urlWithoutQuery.split("/").pop();
+            if (!fileName) {
+                throw new Error(`Invalid model URL: ${pathOrUrl}`);
             }
-            return result;
-        },
-        sendMessage: async (...args) => {
-            const result = await native.sendMessage(...args);
-            recordMemorySnapshot();
-            return result;
-        },
-        sendMessageAsync: (...args) => {
-            const [message, onToken] = args;
-            native.sendMessageAsync(message, (token, done) => {
-                onToken(token, done);
-                if (done) {
-                    recordMemorySnapshot();
-                }
+            console.log(`Checking model at ${pathOrUrl}...`);
+            modelPath = await native.downloadModel(pathOrUrl, fileName, (progress) => {
+                onDownloadProgress?.(progress);
             });
-        },
-        sendMessageWithImage: async (...args) => {
-            const result = await native.sendMessageWithImage(...args);
-            recordMemorySnapshot();
-            return result;
-        },
-        sendMessageWithAudio: async (...args) => {
-            const result = await native.sendMessageWithAudio(...args);
-            recordMemorySnapshot();
-            return result;
-        },
-        getHistory: native.getHistory.bind(native),
-        resetConversation: () => {
-            native.resetConversation();
-            // KV cache is cleared on reset, record the drop
+            console.log(`Model downloaded to: ${modelPath}`);
+        }
+        const result = await native.loadModel(modelPath, config);
+        // Record initial memory snapshot after model load
+        if (tracker) {
+            tracker.reset();
             recordMemorySnapshot();
-        },
-        isReady: native.isReady.bind(native),
-        getStats: native.getStats.bind(native),
-        getMemoryUsage: native.getMemoryUsage.bind(native),
-        close: native.close.bind(native),
-        downloadModel: native.downloadModel.bind(native),
-        deleteModel: native.deleteModel.bind(native),
+        }
+        return result;
     };
+    const SNAPSHOT_TRIGGERS = new Set([
+        "sendMessage",
+        "sendMessageWithImage",
+        "sendMessageWithAudio",
+        "resetConversation",
+    ]);
+    return new Proxy(native, {
+        get(target, prop, receiver) {
+            if (prop === "memoryTracker") {
+                return tracker;
+            }
+            if (prop === "loadModel") {
+                return augmentedLoadModel;
+            }
+            const original = Reflect.get(target, prop, receiver);
+            if (typeof original !== "function") {
+                return original;
+            }
+            if (prop === "sendMessageAsync") {
+                return (message, onToken) => {
+                    return original.call(target, message, (token, done) => {
+                        onToken(token, done);
+                        if (done) {
+                            recordMemorySnapshot();
+                        }
+                    });
+                };
+            }
+            if (SNAPSHOT_TRIGGERS.has(prop)) {
+                return async (...args) => {
+                    const result = await original.apply(target, args);
+                    recordMemorySnapshot();
+                    return result;
+                };
+            }
+            return original.bind(target);
+        },
+    });
 }