Implement Tool macro

Author: mattt

Package.resolved

@@ -1,23 +1,50 @@
-// swift-tools-version: 5.7
+// swift-tools-version: 6.0
 // The swift-tools-version declares the minimum version of Swift required to build this package.
 
+import CompilerPluginSupport
 import PackageDescription
 
 let package = Package(
     name: "Ollama",
     platforms: [
-        .macOS(.v13)
+        .macOS(.v14),
+        .macCatalyst(.v14),
+        .iOS(.v17),
+        .watchOS(.v10),
+        .tvOS(.v17),
+        .visionOS(.v1),
     ],
     products: [
         .library(
             name: "Ollama",
             targets: ["Ollama"])
     ],
+    dependencies: [
+        .package(url: "https://github.com/swiftlang/swift-syntax.git", from: "600.0.0-latest")
+    ],
     targets: [
         .target(
-            name: "Ollama"),
+            name: "Ollama",
+            dependencies: ["OllamaMacro"]),
         .testTarget(
             name: "OllamaTests",
             dependencies: ["Ollama"]),
+
+        .macro(
+            name: "OllamaMacro",
+            dependencies: [
+                .product(name: "SwiftSyntax", package: "swift-syntax"),
+                .product(name: "SwiftSyntaxMacros", package: "swift-syntax"),
+                .product(name: "SwiftCompilerPlugin", package: "swift-syntax"),
+            ]
+        ),
+        .testTarget(
+            name: "OllamaMacroTests",
+            dependencies: [
+                "OllamaMacro",
+                .product(name: "SwiftSyntaxMacroExpansion", package: "swift-syntax"),
+                .product(name: "SwiftSyntaxMacrosGenericTestSupport", package: "swift-syntax"),
+            ]
+        ),
     ]
 )

Package.swift

@@ -10,12 +10,12 @@ import Foundation
 /// allowing you to generate text, chat, create embeddings, and manage models.
 ///
 /// - SeeAlso: [Ollama API Documentation](https://github.com/ollama/ollama/blob/main/docs/api.md)
-open class Client {
+open class Client: @unchecked Sendable {
     /// The default host URL for the Ollama API.
     public static let defaultHost = URL(string: "http://localhost:11434")!
 
     /// A default client instance using the default host.
-    public static let `default` = Client(host: Client.defaultHost)
+    @MainActor public static let `default` = Client(host: Client.defaultHost)
 
     /// The host URL for requests made by the client.
     public let host: URL

Sources/Ollama/Client.swift

@@ -2,7 +2,7 @@ import Foundation
 import RegexBuilder
 
 /// Regex pattern for data URLs
-private let dataURLRegex = Regex {
+nonisolated(unsafe) private let dataURLRegex = Regex {
     "data:"
     Capture {
         ZeroOrMore(.reluctant) {

Sources/Ollama/Extensions/Data+Extensions.swift

@@ -0,0 +1,12 @@
+import Foundation
+
+/// Macro that transforms functions and types into Tools
+@attached(peer, names: prefixed(Tool_))
+public macro Tool() = #externalMacro(module: "OllamaMacro", type: "ToolMacro")
+
+public protocol Tool {
+    associatedtype Input: Codable
+    associatedtype Output: Codable
+
+    static var schema: [String: Value] { get }
+}

Sources/Ollama/Tool.swift

@@ -0,0 +1,157 @@
+import SwiftCompilerPlugin
+import SwiftSyntax
+import SwiftSyntaxBuilder
+import SwiftSyntaxMacros
+
+enum Error: Swift.Error {
+    case unsupportedDeclaration
+}
+
+public struct ToolMacro: PeerMacro {
+    public static func expansion(
+        of node: AttributeSyntax,
+        providingPeersOf declaration: some DeclSyntaxProtocol,
+        in context: some MacroExpansionContext
+    ) throws -> [DeclSyntax] {
+        guard let funcDecl = declaration.as(FunctionDeclSyntax.self) else {
+            throw Error.unsupportedDeclaration
+        }
+
+        let toolName = "Tool_" + funcDecl.name.text
+
+        let inputType = generateInputType(
+            from: funcDecl.signature.parameterClause,
+            toolName: toolName
+        )
+
+        let description = extractDescription(from: funcDecl)
+        let parameterDescriptions = extractParameterDescriptions(from: funcDecl)
+
+        let tool = """
+            enum \(toolName): Tool {
+                \(inputType)
+                typealias Output = \(funcDecl.signature.returnClause?.type.description ?? "Void")
+                
+                static var schema: [String: Value] {
+                    [
+                        "name": "\(funcDecl.name.text)",
+                        "description": "\(description)",
+                        "parameters": \(generateParametersDictionary(from: funcDecl.signature.parameterClause, descriptions: parameterDescriptions))
+                    ]
+                }
+
+                static func call(\(generateCallParameters(from: funcDecl.signature.parameterClause))) throws -> Output {
+                    \(generateFunctionCall(funcDecl: funcDecl))
+                }
+            }
+            """
+
+        let toolDecl: DeclSyntax = "\(raw: tool)"
+        return [toolDecl]
+    }
+
+    private static func generateInputType(
+        from parameters: FunctionParameterClauseSyntax,
+        toolName: String
+    ) -> String {
+        if parameters.parameters.count == 1 {
+            let param = parameters.parameters.first!
+            return "typealias Input = \(param.type)"
+        }
+
+        let structFields = parameters.parameters.map { param in
+            "let \(param.firstName.text): \(param.type)"
+        }.joined(separator: "\n        ")
+
+        let inputStruct = """
+            struct Input: Codable {
+                    \(structFields)
+                }
+            """
+
+        return inputStruct
+    }
+
+    private static func generateCallParameters(from parameters: FunctionParameterClauseSyntax)
+        -> String
+    {
+        return parameters.parameters.map { param in
+            let paramName = param.secondName?.text ?? param.firstName.text
+            return "\(paramName): \(param.type)"
+        }.joined(separator: ", ")
+    }
+
+    private static func generateParametersDictionary(
+        from parameters: FunctionParameterClauseSyntax, descriptions: [String: String]
+    ) -> String {
+        let parameterEntries = parameters.parameters.map { param in
+            let paramName = param.secondName?.text ?? param.firstName.text
+            let paramType = param.type.description
+            if let description = descriptions[paramName] {
+                return """
+                                "\(paramName)": [
+                                    "type": "\(paramType)",
+                                    "description": "\(description)"
+                                ]
+                    """
+            } else {
+                return """
+                                "\(paramName)": [
+                                    "type": "\(paramType)"
+                                ]
+                    """
+            }
+        }.joined(separator: ",\n")
+
+        return "[\n\(parameterEntries)\n        ]"
+    }
+
+    private static func generateFunctionCall(funcDecl: FunctionDeclSyntax) -> String {
+        let params = funcDecl.signature.parameterClause.parameters
+        let arguments = params.map { param in
+            let argName = param.secondName?.text ?? param.firstName.text
+            return "\(param.firstName): \(argName)"
+        }.joined(separator: ", ")
+
+        return "\(funcDecl.name.text)(\(arguments))"
+    }
+
+    private static func extractDescription(from funcDecl: FunctionDeclSyntax) -> String {
+        if let docComment = funcDecl.leadingTrivia.compactMap({ trivia in
+            if case .docLineComment(let comment) = trivia { return comment }
+            return nil
+        }).first {
+            return String(docComment.dropFirst(3).trimmingCharacters(in: .whitespaces))
+        }
+        return "Calls the \(funcDecl.name.text) function"
+    }
+
+    private static func extractParameterDescriptions(from funcDecl: FunctionDeclSyntax) -> [String:
+        String]
+    {
+        var descriptions: [String: String] = [:]
+        let docComments = funcDecl.leadingTrivia.compactMap({ trivia in
+            if case .docLineComment(let comment) = trivia { return comment }
+            return nil
+        })
+        for comment in docComments {
+            let trimmed = comment.dropFirst(3).trimmingCharacters(in: .whitespaces)
+            if trimmed.starts(with: "- Parameter") {
+                let parts = trimmed.dropFirst(12).split(separator: ":", maxSplits: 1)
+                if parts.count == 2 {
+                    let paramName = parts[0].trimmingCharacters(in: .whitespaces)
+                    let description = parts[1].trimmingCharacters(in: .whitespaces)
+                    descriptions[paramName] = description
+                }
+            }
+        }
+        return descriptions
+    }
+}
+
+@main
+struct OllamaToolPlugin: CompilerPlugin {
+    let providingMacros: [Macro.Type] = [
+        ToolMacro.self
+    ]
+}

Sources/OllamaMacro/Tool.swift

@@ -0,0 +1,36 @@
+import SwiftSyntaxMacroExpansion
+import SwiftSyntaxMacros
+import SwiftSyntaxMacrosGenericTestSupport
+import Testing
+
+func assertMacroExpansion(
+    _ originalSource: String,
+    expandedSource expectedExpandedSource: String,
+    macros: [String: Macro.Type],
+    fileID: StaticString = #fileID,
+    filePath: StaticString = #filePath,
+    line: UInt = #line,
+    column: UInt = #column
+) {
+    assertMacroExpansion(
+        originalSource,
+        expandedSource: expectedExpandedSource,
+        macroSpecs: macros.mapValues { MacroSpec(type: $0) },
+        indentationWidth: .spaces(4),
+        failureHandler: { spec in
+            Issue.record(
+                "\(spec.message)",
+                sourceLocation: SourceLocation(
+                    fileID: spec.location.fileID,
+                    filePath: spec.location.filePath,
+                    line: spec.location.line,
+                    column: spec.location.column
+                )
+            )
+        },
+        fileID: fileID,
+        filePath: filePath,
+        line: line,
+        column: column
+    )
+}