Skip to content

Commit 744df1b

Browse files
matttmattt
committed
Implement Tool macro
1 parent 13f8640 commit 744df1b

File tree

11 files changed

+476
-7
lines changed

11 files changed

+476
-7
lines changed

‎Package.resolved

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Package.swift

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,50 @@
1-
// swift-tools-version: 5.7
1+
// swift-tools-version: 6.0
22
// The swift-tools-version declares the minimum version of Swift required to build this package.
33

4+
import CompilerPluginSupport
45
import PackageDescription
56

67
let package = Package(
78
name: "Ollama",
89
platforms: [
9-
.macOS(.v13)
10+
.macOS(.v14),
11+
.macCatalyst(.v14),
12+
.iOS(.v17),
13+
.watchOS(.v10),
14+
.tvOS(.v17),
15+
.visionOS(.v1),
1016
],
1117
products: [
1218
.library(
1319
name: "Ollama",
1420
targets: ["Ollama"])
1521
],
22+
dependencies: [
23+
.package(url: "https://github.com/swiftlang/swift-syntax.git", from: "600.0.0-latest")
24+
],
1625
targets: [
1726
.target(
18-
name: "Ollama"),
27+
name: "Ollama",
28+
dependencies: ["OllamaMacro"]),
1929
.testTarget(
2030
name: "OllamaTests",
2131
dependencies: ["Ollama"]),
32+
33+
.macro(
34+
name: "OllamaMacro",
35+
dependencies: [
36+
.product(name: "SwiftSyntax", package: "swift-syntax"),
37+
.product(name: "SwiftSyntaxMacros", package: "swift-syntax"),
38+
.product(name: "SwiftCompilerPlugin", package: "swift-syntax"),
39+
]
40+
),
41+
.testTarget(
42+
name: "OllamaMacroTests",
43+
dependencies: [
44+
"OllamaMacro",
45+
.product(name: "SwiftSyntaxMacroExpansion", package: "swift-syntax"),
46+
.product(name: "SwiftSyntaxMacrosGenericTestSupport", package: "swift-syntax"),
47+
]
48+
),
2249
]
2350
)

‎Sources/Ollama/Client.swift

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,12 @@ import Foundation
1010
/// allowing you to generate text, chat, create embeddings, and manage models.
1111
///
1212
/// - SeeAlso: [Ollama API Documentation](https://github.com/ollama/ollama/blob/main/docs/api.md)
13-
open class Client {
13+
open class Client: @unchecked Sendable {
1414
/// The default host URL for the Ollama API.
1515
public static let defaultHost = URL(string: "http://localhost:11434")!
1616

1717
/// A default client instance using the default host.
18-
public static let `default` = Client(host: Client.defaultHost)
18+
@MainActor public static let `default` = Client(host: Client.defaultHost)
1919

2020
/// The host URL for requests made by the client.
2121
public let host: URL

‎Sources/Ollama/Extensions/Data+Extensions.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import Foundation
22
import RegexBuilder
33

44
/// Regex pattern for data URLs
5-
private let dataURLRegex = Regex {
5+
nonisolated(unsafe) private let dataURLRegex = Regex {
66
"data:"
77
Capture {
88
ZeroOrMore(.reluctant) {

‎Sources/Ollama/Tool.swift

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
import Foundation
2+
3+
/// Macro that transforms functions and types into Tools
4+
@attached(peer, names: prefixed(Tool_))
5+
public macro Tool() = #externalMacro(module: "OllamaMacro", type: "ToolMacro")
6+
7+
public protocol Tool {
8+
associatedtype Input: Codable
9+
associatedtype Output: Codable
10+
11+
static var schema: [String: Value] { get }
12+
}

‎Sources/OllamaMacro/OllamaToolPlugin.swift

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
import SwiftCompilerPlugin
2+
import SwiftSyntaxMacros
3+
4+
@main
5+
struct OllamaToolPlugin: CompilerPlugin {
6+
let providingMacros: [Macro.Type] = [
7+
ToolMacro.self
8+
]
9+
}

‎Sources/OllamaMacro/ToolMacro.swift

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
import SwiftSyntax
2+
import SwiftSyntaxBuilder
3+
import SwiftSyntaxMacros
4+
5+
enum Error: Swift.Error {
6+
case unsupportedDeclaration
7+
}
8+
9+
public struct ToolMacro: PeerMacro {
10+
public static func expansion(
11+
of node: AttributeSyntax,
12+
providingPeersOf declaration: some DeclSyntaxProtocol,
13+
in context: some MacroExpansionContext
14+
) throws -> [DeclSyntax] {
15+
guard let funcDecl = declaration.as(FunctionDeclSyntax.self) else {
16+
throw Error.unsupportedDeclaration
17+
}
18+
19+
let toolName = "Tool_" + funcDecl.name.text
20+
21+
let inputType = generateInputType(
22+
from: funcDecl.signature.parameterClause,
23+
toolName: toolName
24+
)
25+
26+
let description = extractDescription(from: funcDecl)
27+
let parameterDescriptions = extractParameterDescriptions(from: funcDecl)
28+
29+
let tool = """
30+
enum \(toolName): Tool {
31+
\(inputType)
32+
typealias Output = \(funcDecl.signature.returnClause?.type.description ?? "Void")
33+
34+
static var schema: [String: Value] {
35+
[
36+
"name": "\(funcDecl.name.text)",
37+
"description": "\(description)",
38+
"parameters": \(generateParametersDictionary(from: funcDecl.signature.parameterClause, descriptions: parameterDescriptions))
39+
]
40+
}
41+
42+
static func call(\(generateCallParameters(from: funcDecl.signature.parameterClause))) throws -> Output {
43+
\(generateFunctionCall(funcDecl: funcDecl))
44+
}
45+
}
46+
"""
47+
48+
let toolDecl: DeclSyntax = "\(raw: tool)"
49+
return [toolDecl]
50+
}
51+
52+
private static func generateInputType(
53+
from parameters: FunctionParameterClauseSyntax,
54+
toolName: String
55+
) -> String {
56+
if parameters.parameters.count == 1 {
57+
let param = parameters.parameters.first!
58+
return "typealias Input = \(param.type)"
59+
}
60+
61+
let structFields = parameters.parameters.map { param in
62+
"let \(param.firstName.text): \(param.type)"
63+
}.joined(separator: "\n ")
64+
65+
let inputStruct = """
66+
struct Input: Codable {
67+
\(structFields)
68+
}
69+
"""
70+
71+
return inputStruct
72+
}
73+
74+
private static func generateCallParameters(from parameters: FunctionParameterClauseSyntax)
75+
-> String
76+
{
77+
return parameters.parameters.map { param in
78+
let paramName = param.secondName?.text ?? param.firstName.text
79+
return "\(paramName): \(param.type)"
80+
}.joined(separator: ", ")
81+
}
82+
83+
private static func generateParametersDictionary(
84+
from parameters: FunctionParameterClauseSyntax, descriptions: [String: String]
85+
) -> String {
86+
let parameterEntries = parameters.parameters.map { param in
87+
let paramName = param.secondName?.text ?? param.firstName.text
88+
let paramType = param.type.description
89+
if let description = descriptions[paramName] {
90+
return """
91+
"\(paramName)": [
92+
"type": "\(paramType)",
93+
"description": "\(description)"
94+
]
95+
"""
96+
} else {
97+
return """
98+
"\(paramName)": [
99+
"type": "\(paramType)"
100+
]
101+
"""
102+
}
103+
}.joined(separator: ",\n")
104+
105+
return "[\n\(parameterEntries)\n ]"
106+
}
107+
108+
private static func generateFunctionCall(funcDecl: FunctionDeclSyntax) -> String {
109+
let params = funcDecl.signature.parameterClause.parameters
110+
let arguments = params.map { param in
111+
let argName = param.secondName?.text ?? param.firstName.text
112+
return "\(param.firstName.trimmed): \(argName)"
113+
}.joined(separator: ", ")
114+
115+
return "\(funcDecl.name.text)(\(arguments))"
116+
}
117+
118+
private static func extractDescription(from funcDecl: FunctionDeclSyntax) -> String {
119+
if let docComment = funcDecl.leadingTrivia.compactMap({ trivia in
120+
if case .docLineComment(let comment) = trivia { return comment }
121+
return nil
122+
}).first {
123+
return String(docComment.dropFirst(3).trimmingCharacters(in: .whitespaces))
124+
}
125+
return "Calls the \(funcDecl.name.text) function"
126+
}
127+
128+
private static func extractParameterDescriptions(from funcDecl: FunctionDeclSyntax) -> [String:
129+
String]
130+
{
131+
var descriptions: [String: String] = [:]
132+
let docComments = funcDecl.leadingTrivia.compactMap({ trivia in
133+
if case .docLineComment(let comment) = trivia { return comment }
134+
return nil
135+
})
136+
for comment in docComments {
137+
let trimmed = comment.dropFirst(3).trimmingCharacters(in: .whitespaces)
138+
if trimmed.starts(with: "- Parameter") {
139+
let parts = trimmed.dropFirst(12).split(separator: ":", maxSplits: 1)
140+
if parts.count == 2 {
141+
let paramName = parts[0].trimmingCharacters(in: .whitespaces)
142+
let description = parts[1].trimmingCharacters(in: .whitespaces)
143+
descriptions[paramName] = description
144+
}
145+
}
146+
}
147+
return descriptions
148+
}
149+
}

‎Tests/OllamaMacroTests/Extensions/SwiftTesting+Macro.swift

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import SwiftSyntaxMacroExpansion
2+
import SwiftSyntaxMacros
3+
import SwiftSyntaxMacrosGenericTestSupport
4+
import Testing
5+
6+
func assertMacroExpansion(
7+
_ originalSource: String,
8+
expandedSource expectedExpandedSource: String,
9+
macros: [String: Macro.Type],
10+
fileID: StaticString = #fileID,
11+
filePath: StaticString = #filePath,
12+
line: UInt = #line,
13+
column: UInt = #column
14+
) {
15+
assertMacroExpansion(
16+
originalSource,
17+
expandedSource: expectedExpandedSource,
18+
macroSpecs: macros.mapValues { MacroSpec(type: $0) },
19+
indentationWidth: .spaces(4),
20+
failureHandler: { spec in
21+
Issue.record(
22+
"\(spec.message)",
23+
sourceLocation: SourceLocation(
24+
fileID: spec.location.fileID,
25+
filePath: spec.location.filePath,
26+
line: spec.location.line,
27+
column: spec.location.column
28+
)
29+
)
30+
},
31+
fileID: fileID,
32+
filePath: filePath,
33+
line: line,
34+
column: column
35+
)
36+
}

0 commit comments

Comments
 (0)