diff --git a/Libraries/MLXLMCommon/Tool/Parsers/PythonicToolCallParser.swift b/Libraries/MLXLMCommon/Tool/Parsers/PythonicToolCallParser.swift new file mode 100644 index 000000000..73daec340 --- /dev/null +++ b/Libraries/MLXLMCommon/Tool/Parsers/PythonicToolCallParser.swift @@ -0,0 +1,100 @@ +// Copyright © 2025 Apple Inc. + +import Foundation + +/// Parser for Pythonic tool call format: [function_name(arg1='value1', arg2='value2')] +/// Used by LFM2.5 and similar models that output tool calls in Python function call syntax. +/// Reference: LiquidAI LFM2.5 chat template format +public struct PythonicToolCallParser: ToolCallParser, Sendable { + public let startTag: String? + public let endTag: String? + + public init(startTag: String? = nil, endTag: String? = nil) { + self.startTag = startTag + self.endTag = endTag + } + + public func parse(content: String, tools: [[String: any Sendable]]?) -> ToolCall? { + var text = content + + // Strip tags if present + if let start = startTag, let startRange = text.range(of: start) { + text = String(text[startRange.upperBound...]) + } + if let end = endTag, let endRange = text.range(of: end) { + text = String(text[.. [String: any Sendable] { + var arguments: [String: any Sendable] = [:] + + // Pattern for key=value pairs, handling quoted strings with possible commas inside + // This handles: key='value', key="value", key=123, key=True, key=None + let argPattern = #"(\w+)\s*=\s*('(?:[^'\\]|\\.)*'|"(?:[^"\\]|\\.)*"|[^,\)]+)"# + + guard let regex = try? NSRegularExpression(pattern: argPattern, options: []) else { + return arguments + } + + let matches = regex.matches( + in: argsString, options: [], range: NSRange(argsString.startIndex..., in: argsString)) + + for match in matches { + guard let keyRange = Range(match.range(at: 1), in: argsString), + let valueRange = Range(match.range(at: 2), in: argsString) + else { continue } + + let key = String(argsString[keyRange]) + var value = String(argsString[valueRange]).trimmingCharacters(in: .whitespaces) + + // Remove surrounding quotes if present + if (value.hasPrefix("'") && value.hasSuffix("'")) + || (value.hasPrefix("\"") && value.hasSuffix("\"")) + { + value = String(value.dropFirst().dropLast()) + // Unescape escaped quotes + value = value.replacingOccurrences(of: "\\'", with: "'") + value = value.replacingOccurrences(of: "\\\"", with: "\"") + value = value.replacingOccurrences(of: "\\\\", with: "\\") + } + + // Convert value based on schema type if available + arguments[key] = convertParameterValue( + value, paramName: key, funcName: funcName, tools: tools) + } + + return arguments + } +} diff --git a/Libraries/MLXLMCommon/Tool/ToolCallFormat.swift b/Libraries/MLXLMCommon/Tool/ToolCallFormat.swift index 3b39bf608..33a79b03d 100644 --- a/Libraries/MLXLMCommon/Tool/ToolCallFormat.swift +++ b/Libraries/MLXLMCommon/Tool/ToolCallFormat.swift @@ -42,8 +42,8 @@ public enum ToolCallFormat: String, Sendable, Codable, CaseIterable { /// Example: `{"name": "func", "arguments": {...}}` case json - /// LFM2 JSON format with model-specific tags. - /// Example: `<|tool_call_start|>{"name": "func", "arguments": {...}}<|tool_call_end|>` + /// LFM2/LFM2.5 Pythonic format with model-specific tags. + /// Example: `<|tool_call_start|>[func(arg='value')]<|tool_call_end|>` case lfm2 /// XML function format used by Qwen3 Coder. @@ -75,7 +75,8 @@ public enum ToolCallFormat: String, Sendable, Codable, CaseIterable { case .json: return JSONToolCallParser(startTag: "", endTag: "") case .lfm2: - return JSONToolCallParser(startTag: "<|tool_call_start|>", endTag: "<|tool_call_end|>") + return PythonicToolCallParser( + startTag: "<|tool_call_start|>", endTag: "<|tool_call_end|>") case .xmlFunction: return XMLFunctionParser() case .glm4: @@ -97,15 +98,23 @@ public enum ToolCallFormat: String, Sendable, Codable, CaseIterable { /// - Parameter modelType: The `model_type` value from config.json /// - Returns: The appropriate `ToolCallFormat`, or `nil` to use the default format public static func infer(from modelType: String) -> ToolCallFormat? { - switch modelType.lowercased() { - case "lfm2", "lfm2_moe": + let type = modelType.lowercased() + + // LFM2 family (lfm2, lfm2_moe, lfm2_5, lfm25, etc.) + if type.hasPrefix("lfm2") { return .lfm2 - case "glm4", "glm4_moe", "glm4_moe_lite": + } + + // GLM4 family (glm4, glm4_moe, glm4_moe_lite, etc.) + if type.hasPrefix("glm4") { return .glm4 - case "gemma": + } + + // Gemma + if type == "gemma" { return .gemma - default: - return nil } + + return nil } } diff --git a/Tests/MLXLMTests/ToolTests.swift b/Tests/MLXLMTests/ToolTests.swift index b2b312b8b..96f124fd5 100644 --- a/Tests/MLXLMTests/ToolTests.swift +++ b/Tests/MLXLMTests/ToolTests.swift @@ -100,8 +100,8 @@ struct ToolTests { #expect(toolCall.function.arguments["location"] == .string("Paris")) } - @Test("Test JSON Tool Call Parser - LFM2 Tags") - func testJSONParserLFM2Tags() throws { + @Test("Test JSON Tool Call Parser - Custom Tags") + func testJSONParserCustomTags() throws { let parser = JSONToolCallParser( startTag: "<|tool_call_start|>", endTag: "<|tool_call_end|>") let content = @@ -113,11 +113,93 @@ struct ToolTests { #expect(toolCall.function.arguments["query"] == .string("swift programming")) } - @Test("Test LFM2 Format via ToolCallProcessor") + // MARK: - Pythonic Format Tests (LFM2/LFM2.5) + + @Test("Test Pythonic Tool Call Parser - Basic") + func testPythonicParserBasic() throws { + let parser = PythonicToolCallParser( + startTag: "<|tool_call_start|>", endTag: "<|tool_call_end|>") + let content = + "<|tool_call_start|>[get_weather(location='Paris', unit='celsius')]<|tool_call_end|>" + + let toolCall = try #require(parser.parse(content: content, tools: nil)) + + #expect(toolCall.function.name == "get_weather") + #expect(toolCall.function.arguments["location"] == .string("Paris")) + #expect(toolCall.function.arguments["unit"] == .string("celsius")) + } + + @Test("Test Pythonic Tool Call Parser - Double Quotes") + func testPythonicParserDoubleQuotes() throws { + let parser = PythonicToolCallParser( + startTag: "<|tool_call_start|>", endTag: "<|tool_call_end|>") + let content = + "<|tool_call_start|>[search(query=\"swift programming\")]<|tool_call_end|>" + + let toolCall = try #require(parser.parse(content: content, tools: nil)) + + #expect(toolCall.function.name == "search") + #expect(toolCall.function.arguments["query"] == .string("swift programming")) + } + + @Test("Test Pythonic Tool Call Parser - Without Brackets") + func testPythonicParserWithoutBrackets() throws { + let parser = PythonicToolCallParser( + startTag: "<|tool_call_start|>", endTag: "<|tool_call_end|>") + let content = + "<|tool_call_start|>current_time(timezone='UTC')<|tool_call_end|>" + + let toolCall = try #require(parser.parse(content: content, tools: nil)) + + #expect(toolCall.function.name == "current_time") + #expect(toolCall.function.arguments["timezone"] == .string("UTC")) + } + + @Test("Test Pythonic Tool Call Parser - No Arguments") + func testPythonicParserNoArguments() throws { + let parser = PythonicToolCallParser( + startTag: "<|tool_call_start|>", endTag: "<|tool_call_end|>") + let content = + "<|tool_call_start|>[current_time()]<|tool_call_end|>" + + let toolCall = try #require(parser.parse(content: content, tools: nil)) + + #expect(toolCall.function.name == "current_time") + #expect(toolCall.function.arguments.isEmpty) + } + + @Test("Test Pythonic Tool Call Parser - Type Conversion") + func testPythonicParserTypeConversion() throws { + let parser = PythonicToolCallParser( + startTag: "<|tool_call_start|>", endTag: "<|tool_call_end|>") + let tools: [[String: any Sendable]] = [ + [ + "function": [ + "name": "set_temperature", + "parameters": [ + "properties": [ + "value": ["type": "integer"], + "enabled": ["type": "boolean"], + ] + ], + ] as [String: any Sendable] + ] + ] + let content = + "<|tool_call_start|>[set_temperature(value='25', enabled='true')]<|tool_call_end|>" + + let toolCall = try #require(parser.parse(content: content, tools: tools)) + + #expect(toolCall.function.name == "set_temperature") + #expect(toolCall.function.arguments["value"] == .int(25)) + #expect(toolCall.function.arguments["enabled"] == .bool(true)) + } + + @Test("Test LFM2 Format via ToolCallProcessor - Pythonic") func testLFM2FormatProcessor() throws { let processor = ToolCallProcessor(format: .lfm2) let content = - "<|tool_call_start|>{\"name\": \"calculator\", \"arguments\": {\"expression\": \"2+2\"}}<|tool_call_end|>" + "<|tool_call_start|>[calculator(expression='2+2')]<|tool_call_end|>" _ = processor.processChunk(content) @@ -317,15 +399,20 @@ struct ToolTests { @Test("Test ToolCallFormat Inference from Model Type") func testToolCallFormatInference() throws { - // LFM2 models + // LFM2 models (prefix matching) #expect(ToolCallFormat.infer(from: "lfm2") == .lfm2) #expect(ToolCallFormat.infer(from: "LFM2") == .lfm2) #expect(ToolCallFormat.infer(from: "lfm2_moe") == .lfm2) + #expect(ToolCallFormat.infer(from: "lfm2_5") == .lfm2) + #expect(ToolCallFormat.infer(from: "LFM2_5") == .lfm2) + #expect(ToolCallFormat.infer(from: "lfm25") == .lfm2) - // GLM4 models + // GLM4 models (prefix matching) #expect(ToolCallFormat.infer(from: "glm4") == .glm4) #expect(ToolCallFormat.infer(from: "glm4_moe") == .glm4) #expect(ToolCallFormat.infer(from: "glm4_moe_lite") == .glm4) + #expect(ToolCallFormat.infer(from: "glm4_5") == .glm4) + #expect(ToolCallFormat.infer(from: "GLM4_5") == .glm4) // Gemma models #expect(ToolCallFormat.infer(from: "gemma") == .gemma)