react-native-nitro-mlx 0.2.1 → 0.3.0

package/MLXReactNative.podspec

@@ -24,7 +24,7 @@ Pod::Spec.new do |s|
 
   spm_dependency(s,
     url: "https://github.com/ml-explore/mlx-swift-lm.git",
-    requirement: {kind: "upToNextMinorVersion", minimumVersion: "2.29.2"},
+    requirement: {kind: "upToNextMinorVersion", minimumVersion: "2.29.3"},
     products: ["MLXLLM", "MLXLMCommon"]
   )
 

package/ios/Sources/HybridLLM.swift

@@ -3,6 +3,7 @@ import NitroModules
 internal import MLX
 internal import MLXLLM
 internal import MLXLMCommon
+internal import Tokenizers
 
 class HybridLLM: HybridLLMSpec {
     private var session: ChatSession?
@@ -17,6 +18,10 @@ class HybridLLM: HybridLLMSpec {
     private var modelFactory: ModelFactory = LLMModelFactory.shared
     private var manageHistory: Bool = false
     private var messageHistory: [LLMMessage] = []
+    private var loadTask: Task<Void, Error>?
+
+    private var tools: [ToolDefinition] = []
+    private var toolSchemas: [ToolSpec] = []
 
     var isLoaded: Bool { session != nil }
     var isGenerating: Bool { currentTask != nil }
@@ -60,51 +65,95 @@ class HybridLLM: HybridLLMSpec {
                      allocatedMB, cacheMB, peakMB)
     }
 
+    private func buildToolSchema(from tool: ToolDefinition) -> ToolSpec {
+        var properties: [String: [String: Any]] = [:]
+        var required: [String] = []
+
+        for param in tool.parameters {
+            properties[param.name] = [
+                "type": param.type,
+                "description": param.description
+            ]
+            if param.required {
+                required.append(param.name)
+            }
+        }
+
+        return [
+            "type": "function",
+            "function": [
+                "name": tool.name,
+                "description": tool.description,
+                "parameters": [
+                    "type": "object",
+                    "properties": properties,
+                    "required": required
+                ]
+            ]
+        ] as ToolSpec
+    }
+
     func load(modelId: String, options: LLMLoadOptions?) throws -> Promise<Void> {
-        return Promise.async { [self] in
-            MLX.GPU.set(cacheLimit: 2000000)
+        self.loadTask?.cancel()
 
-            self.currentTask?.cancel()
-            self.currentTask = nil
-            self.session = nil
-            self.container = nil
-            MLX.GPU.clearCache()
+        return Promise.async { [self] in
+            let task = Task { @MainActor in
+                MLX.GPU.set(cacheLimit: 2000000)
 
-            let memoryAfterCleanup = self.getMemoryUsage()
-            let gpuAfterCleanup = self.getGPUMemoryUsage()
-            log("After cleanup - Host: \(memoryAfterCleanup), GPU: \(gpuAfterCleanup)")
+                self.currentTask?.cancel()
+                self.currentTask = nil
+                self.session = nil
+                self.container = nil
+                self.tools = []
+                self.toolSchemas = []
+                MLX.GPU.clearCache()
+
+                let memoryAfterCleanup = self.getMemoryUsage()
+                let gpuAfterCleanup = self.getGPUMemoryUsage()
+                log("After cleanup - Host: \(memoryAfterCleanup), GPU: \(gpuAfterCleanup)")
+
+                let modelDir = await ModelDownloader.shared.getModelDirectory(modelId: modelId)
+                log("Loading from directory: \(modelDir.path)")
+
+                let config = ModelConfiguration(directory: modelDir)
+                let loadedContainer = try await self.modelFactory.loadContainer(
+                    configuration: config
+                ) { progress in
+                    options?.onProgress?(progress.fractionCompleted)
+                }
 
-            let modelDir = await ModelDownloader.shared.getModelDirectory(modelId: modelId)
-            log("Loading from directory: \(modelDir.path)")
+                try Task.checkCancellation()
 
-            let config = ModelConfiguration(directory: modelDir)
-            let loadedContainer = try await modelFactory.loadContainer(
-                configuration: config
-            ) { progress in
-                options?.onProgress?(progress.fractionCompleted)
-            }
+                let memoryAfterContainer = self.getMemoryUsage()
+                let gpuAfterContainer = self.getGPUMemoryUsage()
+                log("Model loaded - Host: \(memoryAfterContainer), GPU: \(gpuAfterContainer)")
 
-            let memoryAfterContainer = self.getMemoryUsage()
-            let gpuAfterContainer = self.getGPUMemoryUsage()
-            log("Model loaded - Host: \(memoryAfterContainer), GPU: \(gpuAfterContainer)")
+                if let jsTools = options?.tools {
+                    self.tools = jsTools
+                    self.toolSchemas = jsTools.map { self.buildToolSchema(from: $0) }
+                    log("Loaded \(self.tools.count) tools: \(self.tools.map { $0.name })")
+                }
 
-            // Convert [LLMMessage]? to [String: Any]?
-            let additionalContextDict: [String: Any]? = if let messages = options?.additionalContext {
-                ["messages": messages.map { ["role": $0.role, "content": $0.content] }]
-            } else {
-                nil
-            }
+                let additionalContextDict: [String: Any]? = if let messages = options?.additionalContext {
+                    ["messages": messages.map { ["role": $0.role, "content": $0.content] }]
+                } else {
+                    nil
+                }
 
-            self.container = loadedContainer
-            self.session = ChatSession(loadedContainer, instructions: self.systemPrompt, additionalContext: additionalContextDict)
-            self.modelId = modelId
+                self.container = loadedContainer
+                self.session = ChatSession(loadedContainer, instructions: self.systemPrompt, additionalContext: additionalContextDict)
+                self.modelId = modelId
 
-            self.manageHistory = options?.manageHistory ?? false
-            self.messageHistory = options?.additionalContext ?? []
+                self.manageHistory = options?.manageHistory ?? false
+                self.messageHistory = options?.additionalContext ?? []
 
-            if self.manageHistory {
-                log("History management enabled with \(self.messageHistory.count) initial messages")
+                if self.manageHistory {
+                    log("History management enabled with \(self.messageHistory.count) initial messages")
+                }
             }
+
+            self.loadTask = task
+            try await task.value
         }
     }
 
@@ -143,8 +192,14 @@ class HybridLLM: HybridLLMSpec {
         }
     }
 
-    func stream(prompt: String, onToken: @escaping (String) -> Void) throws -> Promise<String> {
-        guard let session = session else {
+    private let maxToolCallDepth = 10
+
+    func stream(
+        prompt: String,
+        onToken: @escaping (String) -> Void,
+        onToolCall: ((String, String) -> Void)?
+    ) throws -> Promise<String> {
+        guard let container = container as? ModelContainer else {
             throw LLMError.notLoaded
         }
 
@@ -154,22 +209,24 @@ class HybridLLM: HybridLLMSpec {
             }
 
             let task = Task<String, Error> {
-                var result = ""
-                var tokenCount = 0
                 let startTime = Date()
                 var firstTokenTime: Date?
+                var tokenCount = 0
 
-                log("Streaming response for: \(prompt.prefix(50))...")
-                for try await chunk in session.streamResponse(to: prompt) {
-                    if Task.isCancelled { break }
-
-                    if firstTokenTime == nil {
-                        firstTokenTime = Date()
-                    }
-                    tokenCount += 1
-                    result += chunk
-                    onToken(chunk)
-                }
+                let result = try await self.performGeneration(
+                    container: container,
+                    prompt: prompt,
+                    toolResults: nil,
+                    depth: 0,
+                    onToken: { token in
+                        if firstTokenTime == nil {
+                            firstTokenTime = Date()
+                        }
+                        tokenCount += 1
+                        onToken(token)
+                    },
+                    onToolCall: onToolCall ?? { _, _ in }
+                )
 
                 let endTime = Date()
                 let totalTime = endTime.timeIntervalSince(startTime) * 1000
@@ -205,12 +262,200 @@ class HybridLLM: HybridLLMSpec {
         }
     }
 
+    private func performGeneration(
+        container: ModelContainer,
+        prompt: String,
+        toolResults: [String]?,
+        depth: Int,
+        onToken: @escaping (String) -> Void,
+        onToolCall: @escaping (String, String) -> Void
+    ) async throws -> String {
+        if depth >= maxToolCallDepth {
+            log("Max tool call depth reached (\(maxToolCallDepth))")
+            return ""
+        }
+
+        var output = ""
+        var pendingToolCalls: [(tool: ToolDefinition, args: [String: Any], argsJson: String)] = []
+
+        var chat: [Chat.Message] = []
+
+        if !self.systemPrompt.isEmpty {
+            chat.append(.system(self.systemPrompt))
+        }
+
+        for msg in self.messageHistory {
+            switch msg.role {
+            case "user": chat.append(.user(msg.content))
+            case "assistant": chat.append(.assistant(msg.content))
+            case "system": chat.append(.system(msg.content))
+            case "tool": chat.append(.tool(msg.content))
+            default: break
+            }
+        }
+
+        if depth == 0 {
+            chat.append(.user(prompt))
+        }
+
+        if let toolResults = toolResults {
+            for result in toolResults {
+                chat.append(.tool(result))
+            }
+        }
+
+        let userInput = UserInput(
+            chat: chat,
+            tools: !self.toolSchemas.isEmpty ? self.toolSchemas : nil
+        )
+
+        let lmInput = try await container.prepare(input: userInput)
+
+        let stream = try await container.perform { context in
+            let parameters = GenerateParameters(maxTokens: 2048, temperature: 0.7)
+            return try MLXLMCommon.generate(
+                input: lmInput,
+                parameters: parameters,
+                context: context
+            )
+        }
+
+        for await generation in stream {
+            if Task.isCancelled { break }
+
+            switch generation {
+            case .chunk(let text):
+                output += text
+                onToken(text)
+
+            case .toolCall(let toolCall):
+                log("Tool call detected: \(toolCall.function.name)")
+
+                guard let tool = self.tools.first(where: { $0.name == toolCall.function.name }) else {
+                    log("Unknown tool: \(toolCall.function.name)")
+                    continue
+                }
+
+                let argsDict = self.convertToolCallArguments(toolCall.function.arguments)
+                let argsJson = self.dictionaryToJson(argsDict)
+
+                pendingToolCalls.append((tool: tool, args: argsDict, argsJson: argsJson))
+                onToolCall(toolCall.function.name, argsJson)
+
+            case .info(let info):
+                log("Generation info: \(info.generationTokenCount) tokens, \(String(format: "%.1f", info.tokensPerSecond)) tokens/s")
+            }
+        }
+
+        if !pendingToolCalls.isEmpty {
+            log("Executing \(pendingToolCalls.count) tool call(s)")
+
+            var allToolResults: [String] = []
+
+            for (tool, argsDict, _) in pendingToolCalls {
+                do {
+                    let argsAnyMap = self.dictionaryToAnyMap(argsDict)
+                    let outerPromise = tool.handler(argsAnyMap)
+                    let innerPromise = try await outerPromise.await()
+                    let resultAnyMap = try await innerPromise.await()
+                    let resultDict = self.anyMapToDictionary(resultAnyMap)
+                    let resultJson = self.dictionaryToJson(resultDict)
+
+                    log("Tool result for \(tool.name): \(resultJson.prefix(100))...")
+                    allToolResults.append(resultJson)
+                } catch {
+                    log("Tool execution error for \(tool.name): \(error)")
+                    allToolResults.append("{\"error\": \"Tool execution failed\"}")
+                }
+            }
+
+            if !output.isEmpty {
+                self.messageHistory.append(LLMMessage(role: "assistant", content: output))
+            }
+
+            if depth == 0 {
+                self.messageHistory.append(LLMMessage(role: "user", content: prompt))
+            }
+
+            for result in allToolResults {
+                self.messageHistory.append(LLMMessage(role: "tool", content: result))
+            }
+
+            onToken("\u{200B}")
+
+            let continuation = try await self.performGeneration(
+                container: container,
+                prompt: prompt,
+                toolResults: allToolResults,
+                depth: depth + 1,
+                onToken: onToken,
+                onToolCall: onToolCall
+            )
+
+            return output + continuation
+        }
+
+        return output
+    }
+
+    private func convertToolCallArguments(_ arguments: [String: JSONValue]) -> [String: Any] {
+        var result: [String: Any] = [:]
+        for (key, value) in arguments {
+            result[key] = value.anyValue
+        }
+        return result
+    }
+
+    private func dictionaryToJson(_ dict: [String: Any]) -> String {
+        guard let data = try? JSONSerialization.data(withJSONObject: dict),
+              let json = String(data: data, encoding: .utf8) else {
+            return "{}"
+        }
+        return json
+    }
+
+    private func dictionaryToAnyMap(_ dict: [String: Any]) -> AnyMap {
+        let anyMap = AnyMap()
+        for (key, value) in dict {
+            switch value {
+            case let stringValue as String:
+                anyMap.setString(key: key, value: stringValue)
+            case let doubleValue as Double:
+                anyMap.setDouble(key: key, value: doubleValue)
+            case let intValue as Int:
+                anyMap.setDouble(key: key, value: Double(intValue))
+            case let boolValue as Bool:
+                anyMap.setBoolean(key: key, value: boolValue)
+            default:
+                anyMap.setString(key: key, value: String(describing: value))
+            }
+        }
+        return anyMap
+    }
+
+    private func anyMapToDictionary(_ anyMap: AnyMap) -> [String: Any] {
+        var dict: [String: Any] = [:]
+        for key in anyMap.getAllKeys() {
+            if anyMap.isString(key: key) {
+                dict[key] = anyMap.getString(key: key)
+            } else if anyMap.isDouble(key: key) {
+                dict[key] = anyMap.getDouble(key: key)
+            } else if anyMap.isBool(key: key) {
+                dict[key] = anyMap.getBoolean(key: key)
+            }
+        }
+        return dict
+    }
+
     func stop() throws {
         currentTask?.cancel()
         currentTask = nil
     }
 
     func unload() throws {
+        loadTask?.cancel()
+        loadTask = nil
+
         let memoryBefore = getMemoryUsage()
         let gpuBefore = getGPUMemoryUsage()
         log("Before unload - Host: \(memoryBefore), GPU: \(gpuBefore)")
@@ -219,6 +464,8 @@ class HybridLLM: HybridLLMSpec {
         currentTask = nil
         session = nil
         container = nil
+        tools = []
+        toolSchemas = []
         messageHistory = []
         manageHistory = false
         modelId = ""

package/lib/module/index.js

@@ -3,4 +3,5 @@
 export { LLM } from "./llm.js";
 export { ModelManager } from "./modelManager.js";
 export { MLXModel, MLXModels, ModelFamily, ModelProvider } from "./models.js";
+export { createTool } from "./tool-utils.js";
 //# sourceMappingURL=index.js.map

package/lib/module/index.js.map

@@ -1 +1 @@
-{"version":3,"names":["LLM","ModelManager","MLXModel","MLXModels","ModelFamily","ModelProvider"],"sourceRoot":"../../src","sources":["index.ts"],"mappings":";;AAAA,SAASA,GAAG,QAAsB,UAAO;AACzC,SAASC,YAAY,QAAQ,mBAAgB;AAC7C,SACEC,QAAQ,EACRC,SAAS,EACTC,WAAW,EAEXC,aAAa,QAER,aAAU","ignoreList":[]}
+{"version":3,"names":["LLM","ModelManager","MLXModel","MLXModels","ModelFamily","ModelProvider","createTool"],"sourceRoot":"../../src","sources":["index.ts"],"mappings":";;AAAA,SAASA,GAAG,QAA8D,UAAO;AACjF,SAASC,YAAY,QAAQ,mBAAgB;AAC7C,SACEC,QAAQ,EACRC,SAAS,EACTC,WAAW,EAEXC,aAAa,QAER,aAAU;AAUjB,SAASC,UAAU,QAAqC,iBAAc","ignoreList":[]}

package/lib/module/llm.js

@@ -50,13 +50,43 @@ export const LLM = {
     return getInstance().generate(prompt);
   },
   /**
-   * Stream a response token by token.
+   * Stream a response token by token with optional tool calling support.
+   * Tools must be provided when loading the model via `load()` options.
+   * Tools are automatically executed when the model calls them.
    * @param prompt - The input text to generate a response for
    * @param onToken - Callback invoked for each generated token
+   * @param onToolCall - Optional callback invoked when a tool is called.
+   *   Receives the current tool call and an accumulated array of all tool calls so far.
    * @returns The complete generated text
    */
-  stream(prompt, onToken) {
-    return getInstance().stream(prompt, onToken);
+  stream(prompt, onToken, onToolCall) {
+    const accumulatedToolCalls = [];
+    return getInstance().stream(prompt, onToken, (name, argsJson) => {
+      if (onToolCall) {
+        try {
+          const args = JSON.parse(argsJson);
+          const toolCall = {
+            name,
+            arguments: args
+          };
+          accumulatedToolCalls.push(toolCall);
+          onToolCall({
+            toolCall,
+            allToolCalls: [...accumulatedToolCalls]
+          });
+        } catch {
+          const toolCall = {
+            name,
+            arguments: {}
+          };
+          accumulatedToolCalls.push(toolCall);
+          onToolCall({
+            toolCall,
+            allToolCalls: [...accumulatedToolCalls]
+          });
+        }
+      }
+    });
   },
   /**
    * Stop the current generation. Safe to call even if not generating.

package/lib/module/llm.js.map

@@ -1 +1 @@
-{"version":3,"names":["NitroModules","instance","getInstance","createHybridObject","LLM","load","modelId","options","generate","prompt","stream","onToken","stop","unload","getLastGenerationStats","getHistory","clearHistory","isLoaded","isGenerating","debug","value","systemPrompt"],"sourceRoot":"../../src","sources":["llm.ts"],"mappings":";;AAAA,SAASA,YAAY,QAAQ,4BAA4B;AAGzD,IAAIC,QAAwB,GAAG,IAAI;AAOnC,SAASC,WAAWA,CAAA,EAAY;EAC9B,IAAI,CAACD,QAAQ,EAAE;IACbA,QAAQ,GAAGD,YAAY,CAACG,kBAAkB,CAAU,KAAK,CAAC;EAC5D;EACA,OAAOF,QAAQ;AACjB;;AAEA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA,OAAO,MAAMG,GAAG,GAAG;EACjB;AACF;AACA;AACA;AACA;EACEC,IAAIA,CAACC,OAAe,EAAEC,OAAuB,EAAiB;IAC5D,OAAOL,WAAW,CAAC,CAAC,CAACG,IAAI,CAACC,OAAO,EAAEC,OAAO,CAAC;EAC7C,CAAC;EAED;AACF;AACA;AACA;AACA;AACA;EACEC,QAAQA,CAACC,MAAc,EAAmB;IACxC,OAAOP,WAAW,CAAC,CAAC,CAACM,QAAQ,CAACC,MAAM,CAAC;EACvC,CAAC;EAED;AACF;AACA;AACA;AACA;AACA;EACEC,MAAMA,CAACD,MAAc,EAAEE,OAAgC,EAAmB;IACxE,OAAOT,WAAW,CAAC,CAAC,CAACQ,MAAM,CAACD,MAAM,EAAEE,OAAO,CAAC;EAC9C,CAAC;EAED;AACF;AACA;EACEC,IAAIA,CAAA,EAAS;IACXV,WAAW,CAAC,CAAC,CAACU,IAAI,CAAC,CAAC;EACtB,CAAC;EAED;AACF;AACA;AACA;EACEC,MAAMA,CAAA,EAAS;IACbX,WAAW,CAAC,CAAC,CAACW,MAAM,CAAC,CAAC;EACxB,CAAC;EAED;AACF;AACA;AACA;EACEC,sBAAsBA,CAAA,EAAoB;IACxC,OAAOZ,WAAW,CAAC,CAAC,CAACY,sBAAsB,CAAC,CAAC;EAC/C,CAAC;EAED;AACF;AACA;AACA;EACEC,UAAUA,CAAA,EAAc;IACtB,OAAOb,WAAW,CAAC,CAAC,CAACa,UAAU,CAAC,CAAC;EACnC,CAAC;EAED;AACF;AACA;EACEC,YAAYA,CAAA,EAAS;IACnBd,WAAW,CAAC,CAAC,CAACc,YAAY,CAAC,CAAC;EAC9B,CAAC;EAED;EACA,IAAIC,QAAQA,CAAA,EAAY;IACtB,OAAOf,WAAW,CAAC,CAAC,CAACe,QAAQ;EAC/B,CAAC;EAED;EACA,IAAIC,YAAYA,CAAA,EAAY;IAC1B,OAAOhB,WAAW,CAAC,CAAC,CAACgB,YAAY;EACnC,CAAC;EAED;EACA,IAAIZ,OAAOA,CAAA,EAAW;IACpB,OAAOJ,WAAW,CAAC,CAAC,CAACI,OAAO;EAC9B,CAAC;EAED;EACA,IAAIa,KAAKA,CAAA,EAAY;IACnB,OAAOjB,WAAW,CAAC,CAAC,CAACiB,KAAK;EAC5B,CAAC;EAED,IAAIA,KAAKA,CAACC,KAAc,EAAE;IACxBlB,WAAW,CAAC,CAAC,CAACiB,KAAK,GAAGC,KAAK;EAC7B,CAAC;EAED;AACF;AACA;AACA;AACA;EACE,IAAIC,YAAYA,CAAA,EAAW;IACzB,OAAOnB,WAAW,CAAC,CAAC,CAACmB,YAAY;EACnC,CAAC;EAED,IAAIA,YAAYA,CAACD,KAAa,EAAE;IAC9BlB,WAAW,CAAC,CAAC,CAACmB,YAAY,GAAGD,KAAK;EACpC;AACF,CAAC","ignoreList":[]}
+{"version":3,"names":["NitroModules","instance","getInstance","createHybridObject","LLM","load","modelId","options","generate","prompt","stream","onToken","onToolCall","accumulatedToolCalls","name","argsJson","args","JSON","parse","toolCall","arguments","push","allToolCalls","stop","unload","getLastGenerationStats","getHistory","clearHistory","isLoaded","isGenerating","debug","value","systemPrompt"],"sourceRoot":"../../src","sources":["llm.ts"],"mappings":";;AAAA,SAASA,YAAY,QAAQ,4BAA4B;AAGzD,IAAIC,QAAwB,GAAG,IAAI;AAiBnC,SAASC,WAAWA,CAAA,EAAY;EAC9B,IAAI,CAACD,QAAQ,EAAE;IACbA,QAAQ,GAAGD,YAAY,CAACG,kBAAkB,CAAU,KAAK,CAAC;EAC5D;EACA,OAAOF,QAAQ;AACjB;;AAEA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA,OAAO,MAAMG,GAAG,GAAG;EACjB;AACF;AACA;AACA;AACA;EACEC,IAAIA,CAACC,OAAe,EAAEC,OAAuB,EAAiB;IAC5D,OAAOL,WAAW,CAAC,CAAC,CAACG,IAAI,CAACC,OAAO,EAAEC,OAAO,CAAC;EAC7C,CAAC;EAED;AACF;AACA;AACA;AACA;AACA;EACEC,QAAQA,CAACC,MAAc,EAAmB;IACxC,OAAOP,WAAW,CAAC,CAAC,CAACM,QAAQ,CAACC,MAAM,CAAC;EACvC,CAAC;EAED;AACF;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;EACEC,MAAMA,CACJD,MAAc,EACdE,OAAgC,EAChCC,UAA6C,EAC5B;IACjB,MAAMC,oBAAoC,GAAG,EAAE;IAE/C,OAAOX,WAAW,CAAC,CAAC,CAACQ,MAAM,CAACD,MAAM,EAAEE,OAAO,EAAE,CAACG,IAAY,EAAEC,QAAgB,KAAK;MAC/E,IAAIH,UAAU,EAAE;QACd,IAAI;UACF,MAAMI,IAAI,GAAGC,IAAI,CAACC,KAAK,CAACH,QAAQ,CAA4B;UAC5D,MAAMI,QAAQ,GAAG;YAAEL,IAAI;YAAEM,SAAS,EAAEJ;UAAK,CAAC;UAC1CH,oBAAoB,CAACQ,IAAI,CAACF,QAAQ,CAAC;UACnCP,UAAU,CAAC;YACTO,QAAQ;YACRG,YAAY,EAAE,CAAC,GAAGT,oBAAoB;UACxC,CAAC,CAAC;QACJ,CAAC,CAAC,MAAM;UACN,MAAMM,QAAQ,GAAG;YAAEL,IAAI;YAAEM,SAAS,EAAE,CAAC;UAAE,CAAC;UACxCP,oBAAoB,CAACQ,IAAI,CAACF,QAAQ,CAAC;UACnCP,UAAU,CAAC;YACTO,QAAQ;YACRG,YAAY,EAAE,CAAC,GAAGT,oBAAoB;UACxC,CAAC,CAAC;QACJ;MACF;IACF,CAAC,CAAC;EACJ,CAAC;EAED;AACF;AACA;EACEU,IAAIA,CAAA,EAAS;IACXrB,WAAW,CAAC,CAAC,CAACqB,IAAI,CAAC,CAAC;EACtB,CAAC;EAED;AACF;AACA;AACA;EACEC,MAAMA,CAAA,EAAS;IACbtB,WAAW,CAAC,CAAC,CAACsB,MAAM,CAAC,CAAC;EACxB,CAAC;EAED;AACF;AACA;AACA;EACEC,sBAAsBA,CAAA,EAAoB;IACxC,OAAOvB,WAAW,CAAC,CAAC,CAACuB,sBAAsB,CAAC,CAAC;EAC/C,CAAC;EAED;AACF;AACA;AACA;EACEC,UAAUA,CAAA,EAAc;IACtB,OAAOxB,WAAW,CAAC,CAAC,CAACwB,UAAU,CAAC,CAAC;EACnC,CAAC;EAED;AACF;AACA;EACEC,YAAYA,CAAA,EAAS;IACnBzB,WAAW,CAAC,CAAC,CAACyB,YAAY,CAAC,CAAC;EAC9B,CAAC;EAED;EACA,IAAIC,QAAQA,CAAA,EAAY;IACtB,OAAO1B,WAAW,CAAC,CAAC,CAAC0B,QAAQ;EAC/B,CAAC;EAED;EACA,IAAIC,YAAYA,CAAA,EAAY;IAC1B,OAAO3B,WAAW,CAAC,CAAC,CAAC2B,YAAY;EACnC,CAAC;EAED;EACA,IAAIvB,OAAOA,CAAA,EAAW;IACpB,OAAOJ,WAAW,CAAC,CAAC,CAACI,OAAO;EAC9B,CAAC;EAED;EACA,IAAIwB,KAAKA,CAAA,EAAY;IACnB,OAAO5B,WAAW,CAAC,CAAC,CAAC4B,KAAK;EAC5B,CAAC;EAED,IAAIA,KAAKA,CAACC,KAAc,EAAE;IACxB7B,WAAW,CAAC,CAAC,CAAC4B,KAAK,GAAGC,KAAK;EAC7B,CAAC;EAED;AACF;AACA;AACA;AACA;EACE,IAAIC,YAAYA,CAAA,EAAW;IACzB,OAAO9B,WAAW,CAAC,CAAC,CAAC8B,YAAY;EACnC,CAAC;EAED,IAAIA,YAAYA,CAACD,KAAa,EAAE;IAC9B7B,WAAW,CAAC,CAAC,CAAC8B,YAAY,GAAGD,KAAK;EACpC;AACF,CAAC","ignoreList":[]}

package/lib/module/tool-utils.js

@@ -0,0 +1,56 @@
+"use strict";
+
+function getZodTypeString(zodType) {
+  const typeName = zodType._zod.def.type;
+  switch (typeName) {
+    case 'string':
+      return 'string';
+    case 'number':
+    case 'int':
+      return 'number';
+    case 'boolean':
+      return 'boolean';
+    case 'array':
+      return 'array';
+    case 'object':
+      return 'object';
+    case 'optional':
+      return getZodTypeString(zodType._zod.def.innerType);
+    case 'default':
+      return getZodTypeString(zodType._zod.def.innerType);
+    default:
+      return 'string';
+  }
+}
+function isZodOptional(zodType) {
+  const typeName = zodType._zod.def.type;
+  return typeName === 'optional' || typeName === 'default';
+}
+function zodSchemaToParameters(schema) {
+  const shape = schema._zod.def.shape;
+  const parameters = [];
+  for (const [key, zodType] of Object.entries(shape)) {
+    const zType = zodType;
+    parameters.push({
+      name: key,
+      type: getZodTypeString(zType),
+      description: zType.description ?? '',
+      required: !isZodOptional(zType)
+    });
+  }
+  return parameters;
+}
+export function createTool(definition) {
+  return {
+    name: definition.name,
+    description: definition.description,
+    parameters: zodSchemaToParameters(definition.arguments),
+    handler: async args => {
+      const argsObj = args;
+      const parsedArgs = definition.arguments.parse(argsObj);
+      const result = await definition.handler(parsedArgs);
+      return result;
+    }
+  };
+}
+//# sourceMappingURL=tool-utils.js.map

package/lib/module/tool-utils.js.map

@@ -0,0 +1 @@
+{"version":3,"names":["getZodTypeString","zodType","typeName","_zod","def","type","innerType","isZodOptional","zodSchemaToParameters","schema","shape","parameters","key","Object","entries","zType","push","name","description","required","createTool","definition","arguments","handler","args","argsObj","parsedArgs","parse","result"],"sourceRoot":"../../src","sources":["tool-utils.ts"],"mappings":";;AAcA,SAASA,gBAAgBA,CAACC,OAAkB,EAAqB;EAC/D,MAAMC,QAAQ,GAAGD,OAAO,CAACE,IAAI,CAACC,GAAG,CAACC,IAAI;EACtC,QAAQH,QAAQ;IACd,KAAK,QAAQ;MACX,OAAO,QAAQ;IACjB,KAAK,QAAQ;IACb,KAAK,KAAK;MACR,OAAO,QAAQ;IACjB,KAAK,SAAS;MACZ,OAAO,SAAS;IAClB,KAAK,OAAO;MACV,OAAO,OAAO;IAChB,KAAK,QAAQ;MACX,OAAO,QAAQ;IACjB,KAAK,UAAU;MACb,OAAOF,gBAAgB,CAAEC,OAAO,CAA8BE,IAAI,CAACC,GAAG,CAACE,SAAS,CAAC;IACnF,KAAK,SAAS;MACZ,OAAON,gBAAgB,CAAEC,OAAO,CAA6BE,IAAI,CAACC,GAAG,CAACE,SAAS,CAAC;IAClF;MACE,OAAO,QAAQ;EACnB;AACF;AAEA,SAASC,aAAaA,CAACN,OAAkB,EAAW;EAClD,MAAMC,QAAQ,GAAGD,OAAO,CAACE,IAAI,CAACC,GAAG,CAACC,IAAI;EACtC,OAAOH,QAAQ,KAAK,UAAU,IAAIA,QAAQ,KAAK,SAAS;AAC1D;AAEA,SAASM,qBAAqBA,CAACC,MAAuB,EAAmB;EACvE,MAAMC,KAAK,GAAGD,MAAM,CAACN,IAAI,CAACC,GAAG,CAACM,KAAK;EACnC,MAAMC,UAA2B,GAAG,EAAE;EAEtC,KAAK,MAAM,CAACC,GAAG,EAAEX,OAAO,CAAC,IAAIY,MAAM,CAACC,OAAO,CAACJ,KAAK,CAAC,EAAE;IAClD,MAAMK,KAAK,GAAGd,OAAoB;IAClCU,UAAU,CAACK,IAAI,CAAC;MACdC,IAAI,EAAEL,GAAG;MACTP,IAAI,EAAEL,gBAAgB,CAACe,KAAK,CAAC;MAC7BG,WAAW,EAAEH,KAAK,CAACG,WAAW,IAAI,EAAE;MACpCC,QAAQ,EAAE,CAACZ,aAAa,CAACQ,KAAK;IAChC,CAAC,CAAC;EACJ;EAEA,OAAOJ,UAAU;AACnB;AAEA,OAAO,SAASS,UAAUA,CACxBC,UAAqC,EACrB;EAChB,OAAO;IACLJ,IAAI,EAAEI,UAAU,CAACJ,IAAI;IACrBC,WAAW,EAAEG,UAAU,CAACH,WAAW;IACnCP,UAAU,EAAEH,qBAAqB,CAACa,UAAU,CAACC,SAAS,CAAC;IACvDC,OAAO,EAAE,MAAOC,IAAY,IAAK;MAC/B,MAAMC,OAAO,GAAGD,IAA0C;MAC1D,MAAME,UAAU,GAAGL,UAAU,CAACC,SAAS,CAACK,KAAK,CAACF,OAAO,CAAC;MACtD,MAAMG,MAAM,GAAG,MAAMP,UAAU,CAACE,OAAO,CAACG,UAAU,CAAC;MACnD,OAAOE,MAAM;IACf;EACF,CAAC;AACH","ignoreList":[]}

package/lib/typescript/src/index.d.ts

@@ -1,6 +1,7 @@
-export { LLM, type Message } from './llm';
+export { LLM, type Message, type ToolCallInfo, type ToolCallUpdate } from './llm';
 export { ModelManager } from './modelManager';
 export { MLXModel, MLXModels, ModelFamily, type ModelInfo, ModelProvider, type ModelQuantization, } from './models';
-export type { GenerationStats, LLM as LLMSpec, LLMLoadOptions } from './specs/LLM.nitro';
+export type { GenerationStats, LLM as LLMSpec, LLMLoadOptions, ToolDefinition, ToolParameter, ToolParameterType, } from './specs/LLM.nitro';
 export type { ModelManager as ModelManagerSpec } from './specs/ModelManager.nitro';
+export { createTool, type TypeSafeToolDefinition } from './tool-utils';
 //# sourceMappingURL=index.d.ts.map

package/lib/typescript/src/index.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,GAAG,EAAE,KAAK,OAAO,EAAE,MAAM,OAAO,CAAA;AACzC,OAAO,EAAE,YAAY,EAAE,MAAM,gBAAgB,CAAA;AAC7C,OAAO,EACL,QAAQ,EACR,SAAS,EACT,WAAW,EACX,KAAK,SAAS,EACd,aAAa,EACb,KAAK,iBAAiB,GACvB,MAAM,UAAU,CAAA;AAEjB,YAAY,EAAE,eAAe,EAAE,GAAG,IAAI,OAAO,EAAE,cAAc,EAAE,MAAM,mBAAmB,CAAA;AACxF,YAAY,EAAE,YAAY,IAAI,gBAAgB,EAAE,MAAM,4BAA4B,CAAA"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,GAAG,EAAE,KAAK,OAAO,EAAE,KAAK,YAAY,EAAE,KAAK,cAAc,EAAE,MAAM,OAAO,CAAA;AACjF,OAAO,EAAE,YAAY,EAAE,MAAM,gBAAgB,CAAA;AAC7C,OAAO,EACL,QAAQ,EACR,SAAS,EACT,WAAW,EACX,KAAK,SAAS,EACd,aAAa,EACb,KAAK,iBAAiB,GACvB,MAAM,UAAU,CAAA;AACjB,YAAY,EACV,eAAe,EACf,GAAG,IAAI,OAAO,EACd,cAAc,EACd,cAAc,EACd,aAAa,EACb,iBAAiB,GAClB,MAAM,mBAAmB,CAAA;AAC1B,YAAY,EAAE,YAAY,IAAI,gBAAgB,EAAE,MAAM,4BAA4B,CAAA;AAClF,OAAO,EAAE,UAAU,EAAE,KAAK,sBAAsB,EAAE,MAAM,cAAc,CAAA"}

package/lib/typescript/src/llm.d.ts

@@ -3,6 +3,14 @@ export type Message = {
     role: 'user' | 'assistant' | 'system';
     content: string;
 };
+export type ToolCallInfo = {
+    name: string;
+    arguments: Record<string, unknown>;
+};
+export type ToolCallUpdate = {
+    toolCall: ToolCallInfo;
+    allToolCalls: ToolCallInfo[];
+};
 /**
  * LLM text generation using MLX on Apple Silicon.
  *
@@ -40,12 +48,16 @@ export declare const LLM: {
      */
     generate(prompt: string): Promise<string>;
     /**
-     * Stream a response token by token.
+     * Stream a response token by token with optional tool calling support.
+     * Tools must be provided when loading the model via `load()` options.
+     * Tools are automatically executed when the model calls them.
      * @param prompt - The input text to generate a response for
      * @param onToken - Callback invoked for each generated token
+     * @param onToolCall - Optional callback invoked when a tool is called.
+     *   Receives the current tool call and an accumulated array of all tool calls so far.
      * @returns The complete generated text
      */
-    stream(prompt: string, onToken: (token: string) => void): Promise<string>;
+    stream(prompt: string, onToken: (token: string) => void, onToolCall?: (update: ToolCallUpdate) => void): Promise<string>;
     /**
      * Stop the current generation. Safe to call even if not generating.
      */

package/lib/typescript/src/llm.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"llm.d.ts","sourceRoot":"","sources":["../../../src/llm.ts"],"names":[],"mappings":"AACA,OAAO,KAAK,EAAE,eAAe,EAAE,cAAc,EAAkB,MAAM,mBAAmB,CAAA;AAIxF,MAAM,MAAM,OAAO,GAAG;IACpB,IAAI,EAAE,MAAM,GAAG,WAAW,GAAG,QAAQ,CAAA;IACrC,OAAO,EAAE,MAAM,CAAA;CAChB,CAAA;AASD;;;;;;;;;;;;;;;;;;;;;GAqBG;AACH,eAAO,MAAM,GAAG;IACd;;;;OAIG;kBACW,MAAM,WAAW,cAAc,GAAG,OAAO,CAAC,IAAI,CAAC;IAI7D;;;;;OAKG;qBACc,MAAM,GAAG,OAAO,CAAC,MAAM,CAAC;IAIzC;;;;;OAKG;mBACY,MAAM,WAAW,CAAC,KAAK,EAAE,MAAM,KAAK,IAAI,GAAG,OAAO,CAAC,MAAM,CAAC;IAIzE;;OAEG;YACK,IAAI;IAIZ;;;OAGG;cACO,IAAI;IAId;;;OAGG;8BACuB,eAAe;IAIzC;;;OAGG;kBACW,OAAO,EAAE;IAIvB;;OAEG;oBACa,IAAI;IAIpB,mEAAmE;uBACnD,OAAO;IAIvB,gDAAgD;2BAC5B,OAAO;IAI3B,oEAAoE;sBACrD,MAAM;IAIrB,sCAAsC;WACzB,OAAO;IAQpB;;;;OAIG;kBACiB,MAAM;CAO3B,CAAA"}
+{"version":3,"file":"llm.d.ts","sourceRoot":"","sources":["../../../src/llm.ts"],"names":[],"mappings":"AACA,OAAO,KAAK,EAAE,eAAe,EAAE,cAAc,EAAkB,MAAM,mBAAmB,CAAA;AAIxF,MAAM,MAAM,OAAO,GAAG;IACpB,IAAI,EAAE,MAAM,GAAG,WAAW,GAAG,QAAQ,CAAA;IACrC,OAAO,EAAE,MAAM,CAAA;CAChB,CAAA;AAED,MAAM,MAAM,YAAY,GAAG;IACzB,IAAI,EAAE,MAAM,CAAA;IACZ,SAAS,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAA;CACnC,CAAA;AAED,MAAM,MAAM,cAAc,GAAG;IAC3B,QAAQ,EAAE,YAAY,CAAA;IACtB,YAAY,EAAE,YAAY,EAAE,CAAA;CAC7B,CAAA;AASD;;;;;;;;;;;;;;;;;;;;;GAqBG;AACH,eAAO,MAAM,GAAG;IACd;;;;OAIG;kBACW,MAAM,WAAW,cAAc,GAAG,OAAO,CAAC,IAAI,CAAC;IAI7D;;;;;OAKG;qBACc,MAAM,GAAG,OAAO,CAAC,MAAM,CAAC;IAIzC;;;;;;;;;OASG;mBAEO,MAAM,WACL,CAAC,KAAK,EAAE,MAAM,KAAK,IAAI,eACnB,CAAC,MAAM,EAAE,cAAc,KAAK,IAAI,GAC5C,OAAO,CAAC,MAAM,CAAC;IAyBlB;;OAEG;YACK,IAAI;IAIZ;;;OAGG;cACO,IAAI;IAId;;;OAGG;8BACuB,eAAe;IAIzC;;;OAGG;kBACW,OAAO,EAAE;IAIvB;;OAEG;oBACa,IAAI;IAIpB,mEAAmE;uBACnD,OAAO;IAIvB,gDAAgD;2BAC5B,OAAO;IAI3B,oEAAoE;sBACrD,MAAM;IAIrB,sCAAsC;WACzB,OAAO;IAQpB;;;;OAIG;kBACiB,MAAM;CAO3B,CAAA"}