Separate LlamaKitMacros; Add @Tool Macro; Fix extern LLAMA settings

Package.swift

@@ -31,8 +31,8 @@ var cSettings: [CSetting] =  [
     // NOTE: NEW_LAPACK will required iOS version 16.4+
     // We should consider add this in the future when we drop support for iOS 14
     // (ref: ref: https://developer.apple.com/documentation/accelerate/1513264-cblas_sgemm?language=objc)
-     .define("ACCELERATE_NEW_LAPACK"),
+    .define("ACCELERATE_NEW_LAPACK"),
-     .define("ACCELERATE_LAPACK_ILP64"),
+    .define("ACCELERATE_LAPACK_ILP64")
 ]
 
 #if canImport(Darwin)
@@ -114,6 +114,15 @@ let package = Package(
             ], 
             path: "swift/JSONSchemaMacros"
         ),
+        .macro(
+            name: "LlamaKitMacros",
+            dependencies: [
+                .product(name: "SwiftSyntax", package: "swift-syntax"),
+                .product(name: "SwiftSyntaxMacros", package: "swift-syntax"),
+                .product(name: "SwiftCompilerPlugin", package: "swift-syntax"),
+            ],
+            path: "swift/LlamaKitMacros"
+        ),
         .target(
             name: "JSONSchema",
             dependencies: ["JSONSchemaMacros"],
@@ -121,7 +130,7 @@ let package = Package(
         ),
         .target(
             name: "LlamaKit",
-            dependencies: ["JSONSchema", "LlamaObjC"],
+            dependencies: ["JSONSchema", "LlamaObjC", "LlamaKitMacros"],
             path: "swift/LlamaKit"
         ),
         .testTarget(name: "LlamaKitTests",

common/common.h

@@ -34,10 +34,10 @@ struct common_lora_adapter_container : common_lora_adapter_info {
 };
 
 // build info
-static int LLAMA_BUILD_NUMBER = 0;
+extern int LLAMA_BUILD_NUMBER;
-static char const * LLAMA_COMMIT = "";
+extern char const * LLAMA_COMMIT;
-static char const * LLAMA_COMPILER = "";
+extern char const * LLAMA_COMPILER;
-static char const * LLAMA_BUILD_TARGET = "";
+extern char const * LLAMA_BUILD_TARGET;
 
 struct common_control_vector_load_info;
 

objc/GPTParams.mm

@@ -716,11 +716,4 @@
     gpt_params.input_suffix = [inputSuffix cStringUsingEncoding:NSUTF8StringEncoding];
 }
 
-
-- (LlamaContextParams *)llamaContextParams {
-}
-
-- (LlamaModelParams *)llamaModelParams {
-}
-
 @end

objc/include/GPTParams.h

@@ -256,8 +256,6 @@ typedef NS_ENUM(NSInteger, GGMLSchedPriority) {
 @property (nonatomic, assign) BOOL inputPrefixBOS; // prefix BOS to user inputs, preceding input_prefix
 @property (nonatomic, assign) BOOL ctxShift; // context shift on inifinite text generation
 @property (nonatomic, assign) BOOL displayPrompt; // print prompt before generation
-- (LlamaModelParams *)llamaModelParams;
-- (LlamaContextParams *)llamaContextParams;
 
 @end
 

objc/include/LlamaContext.h

@@ -10,11 +10,11 @@ typedef int32_t LlamaToken;
 
 @interface LlamaContext : NSObject
 
+- (NSUInteger)nCtx;
+
 - (void)attachThreadpool:(GGMLThreadpool *)threadpool
          threadpoolBatch:(GGMLThreadpool *)threadpoolBatch;
 
-- (NSUInteger)nCtx;
-
 // Positive return values does not mean a fatal error, but rather a warning.
 //   0 - success
 //   1 - could not find a KV slot for the batch (try reducing the size of the batch or increase the context)

objc/include/LlamaObjC.h

@@ -2,12 +2,17 @@
 #define LlamaObjC_h
 
 #include <Foundation/Foundation.h>
-#include <llama.h>
-#include <LlamaModel.h>
-#include <LlamaContext.h>
-#include <LlamaSession.h>
 #include <GPTParams.h>
+#include <GPTSampler.h>
+#include <llama.h>
+#include <LlamaBatch.h>
+#include <LlamaContext.h>
+#include <LlamaModel.h>
+#include <LlamaSession.h>
 
-
+int LLAMA_BUILD_NUMBER = 0;
+char const * LLAMA_COMMIT = "unknown";
+char const * LLAMA_COMPILER = "unknown";
+char const * LLAMA_BUILD_TARGET = "unknown";
 
 #endif /* LlamaObjC_h */

swift/JSONSchema/JSONSchema.swift

@@ -1,5 +1,4 @@
 import Foundation
-//import SwiftSyntaxMacros
 
 public struct JSONSchema : Codable {
     public struct Items : Codable {
@@ -71,7 +70,7 @@ public struct _JSONFunctionSchema: Codable {
     public struct Parameters: Codable {
         public let properties: [String: Property]
         public let required: [String]
-        public let type = "object"
+        public var type = "object"
         
         public init(properties: [String : Property], required: [String]) {
             self.properties = properties

swift/JSONSchemaMacros/JSONSchemaMacros.swift

@@ -114,116 +114,9 @@ struct JSONSchemaMacro: ExtensionMacro, MemberMacro {
     }
 }
 
-enum TestError: Error {
-    case message(String)
-}
-
-struct LlamaActorMacro: ExtensionMacro, MemberMacro {
-    static func expansion(of node: AttributeSyntax, providingMembersOf declaration: some DeclGroupSyntax, conformingTo protocols: [TypeSyntax], in context: some MacroExpansionContext) throws -> [DeclSyntax] {
-        [
-            """
-            let session: LlamaToolSession
-            
-            public init(params: GPTParams) async throws {
-                self.session = try await LlamaToolSession(params: params, tools: Self.tools)
-            }
-            """
-        ]
-    }
-    
-    static func expansion(of node: AttributeSyntax,
-                          attachedTo declaration: some DeclGroupSyntax,
-                          providingExtensionsOf type: some TypeSyntaxProtocol,
-                          conformingTo protocols: [TypeSyntax],
-                          in context: some MacroExpansionContext) throws -> [ExtensionDeclSyntax] {
-        var tools: [
-            (name: String,
-             description: String,
-             parameters: [(name: String,
-                           type: String,
-                           description: String)],
-             callableString: String,
-             callableName: String)
-        ] = []
-        for member in declaration.memberBlock.members {
-            let comments = member.leadingTrivia.filter { $0.isComment }
-            
-            guard let member = member.decl.as(FunctionDeclSyntax.self) else {
-                continue
-            }
-            let name = member.name
-            guard case var .docLineComment(description) = comments.first else {
-                throw TestError.message("Missing comment")
-            }
-            description = String(description.dropFirst(3))
-            var parameters: [(name: String, type: String, description: String)] = []
-            var index = 0
-            for parameter in member.signature.parameterClause.parameters {
-                let firstName = parameter.firstName.text
-                let typeName = parameter.type.as(IdentifierTypeSyntax.self)!.name.text
-                guard case var .docLineComment(description) = comments[index + 1] else {
-                    throw TestError.message("Missing comment for \(firstName)")
-                }
-                description = String(description.dropFirst(3))
-                parameters.append((name: firstName, type: typeName, description: description))
-                index += 1
-            }
-            let callableName = context.makeUniqueName(name.text)
-            let callableString = """
-            @dynamicCallable struct \(callableName.text): DynamicCallable {
-                @discardableResult
-                func dynamicallyCall(withKeywordArguments args: [String: Any]) async throws -> String {
-                    \(parameters.map {
-                        "var \($0.name): \($0.type)!"
-                    }.joined(separator: "\n"))
-                    for (key, value) in args {
-                        \(parameters.map {
-                            "if key == \"\($0.name)\" { \($0.name) = value as! \($0.type) }"
-                        }.joined(separator: "\n"))
-                    }
-            
-                    let returnValue = try await \(name.text)(\(parameters.map { "\($0.name): \($0.name)" }.joined(separator: ",")))
-                    let jsonValue = try JSONEncoder().encode(returnValue)
-                    return String(data: jsonValue, encoding: .utf8)!
-                }
-            }
-            """
-            tools.append((name: name.text, description: description,
-                          parameters: parameters,
-                          callableString: callableString,
-                          callableName: callableName.text))
-        }
-        
-        
-        return [
-            .init(extendedType: type,
-                inheritanceClause: .init(inheritedTypes: InheritedTypeListSyntax.init(arrayLiteral: .init(type: IdentifierTypeSyntax(name: "LlamaActor")))),
-                  memberBlock: """
-            {
-                \(raw: tools.map {
-                    $0.callableString
-                }.joined(separator: "\n"))
-            
-                static var tools: [String: (DynamicCallable, _JSONFunctionSchema)] {
-                    [\(raw: tools.map { tool in
-                        """
-                        "\(tool.name)": (\(tool.callableName)(), _JSONFunctionSchema(name: "\(tool.name)", description: "\(tool.description)", parameters: _JSONFunctionSchema.Parameters(properties: \(tool.parameters.count == 0 ? "[:]" : "[" + tool.parameters.map { parameter in
-                            """
-                            "\(parameter.name)": _JSONFunctionSchema.Property(type: \(parameter.type).self, description: "\(parameter.description)"),
-                            """
-                            }.joined() + "]"), required: [])))
-                        """
-                    }.joined(separator: ","))]
-                }
-            }
-            """)
-        ]
-    }
-}
-
 @main
 struct JSONSchemaMacrosPlugin: CompilerPlugin {
     let providingMacros: [Macro.Type] = [
-        JSONSchemaMacro.self, LlamaActorMacro.self
+        JSONSchemaMacro.self
     ]
 }

swift/LlamaKit/LlamaKit.swift

@@ -122,7 +122,7 @@ public actor LlamaToolSession {
         {"name": <function-name>,"arguments": <args-dict>}
         </tool_call>
 
-        Here are the available tools:
+        The first call you will be asked to warm up is to get the user's IP address. Here are the available tools:
         <tools> \(String(data: encoded, encoding: .utf8)!) </tools><|eot_id|>
         """
         params.prompt = prompt
@@ -185,5 +185,9 @@ public extension LlamaActor {
 
 @attached(member, names: arbitrary)
 @attached(extension, conformances: LlamaActor, names: arbitrary)
-public macro llamaActor() = #externalMacro(module: "JSONSchemaMacros",
+public macro llamaActor() = #externalMacro(module: "LlamaKitMacros",
                                            type: "LlamaActorMacro")
+
+@attached(body)
+public macro Tool() = #externalMacro(module: "LlamaKitMacros",
+                                     type: "ToolMacro")

swift/LlamaKitMacros/LlamaKitMacros.swift

@@ -0,0 +1,130 @@
+import Foundation
+import SwiftSyntaxMacros
+import SwiftCompilerPlugin
+import SwiftSyntax
+
+
+enum LlamaKitMacroError: Error {
+    case message(String)
+}
+
+struct ToolMacro: BodyMacro {
+    static func expansion(of node: SwiftSyntax.AttributeSyntax, providingBodyFor declaration: some SwiftSyntax.DeclSyntaxProtocol & SwiftSyntax.WithOptionalCodeBlockSyntax, in context: some SwiftSyntaxMacros.MacroExpansionContext) throws -> [SwiftSyntax.CodeBlockItemSyntax] {
+        []
+    }
+}
+
+struct LlamaActorMacro: ExtensionMacro, MemberMacro {
+    static func expansion(of node: AttributeSyntax, providingMembersOf declaration: some DeclGroupSyntax, conformingTo protocols: [TypeSyntax], in context: some MacroExpansionContext) throws -> [DeclSyntax] {
+        [
+            """
+            let session: LlamaToolSession
+            
+            public init(params: GPTParams) async throws {
+                self.session = try await LlamaToolSession(params: params, tools: Self.tools)
+            }
+            """
+        ]
+    }
+    
+    static func expansion(of node: AttributeSyntax,
+                          attachedTo declaration: some DeclGroupSyntax,
+                          providingExtensionsOf type: some TypeSyntaxProtocol,
+                          conformingTo protocols: [TypeSyntax],
+                          in context: some MacroExpansionContext) throws -> [ExtensionDeclSyntax] {
+        var tools: [
+            (name: String,
+             description: String,
+             parameters: [(name: String,
+                           type: String,
+                           description: String)],
+             callableString: String,
+             callableName: String)
+        ] = []
+        for member in declaration.memberBlock.members {
+            let comments = member.leadingTrivia.filter { $0.isComment }
+            guard let member = member.decl.as(FunctionDeclSyntax.self) else {
+                continue
+            }
+            guard member.attributes.contains(where: { element in
+                element.as(AttributeSyntax.self)?.attributeName.as(IdentifierTypeSyntax.self)?.name.text == "Tool"
+            }) else {
+                continue
+            }
+
+            let name = member.name
+            guard case var .docLineComment(description) = comments.first else {
+                throw LlamaKitMacroError.message("Missing comment")
+            }
+            description = String(description.dropFirst(3))
+            var parameters: [(name: String, type: String, description: String)] = []
+            var index = 0
+            for parameter in member.signature.parameterClause.parameters {
+                let firstName = parameter.firstName.text
+                let typeName = parameter.type.as(IdentifierTypeSyntax.self)!.name.text
+                guard case var .docLineComment(description) = comments[index + 1] else {
+                    throw LlamaKitMacroError.message("Missing comment for \(firstName)")
+                }
+                description = String(description.dropFirst(3))
+                parameters.append((name: firstName, type: typeName, description: description))
+                index += 1
+            }
+            let callableName = context.makeUniqueName(name.text)
+            let callableString = """
+            @dynamicCallable struct \(callableName.text): DynamicCallable {
+                @discardableResult
+                func dynamicallyCall(withKeywordArguments args: [String: Any]) async throws -> String {
+                    \(parameters.map {
+                        "var \($0.name): \($0.type)!"
+                    }.joined(separator: "\n"))
+                    for (key, value) in args {
+                        \(parameters.map {
+                            "if key == \"\($0.name)\" { \($0.name) = value as! \($0.type) }"
+                        }.joined(separator: "\n"))
+                    }
+            
+                    let returnValue = try await \(name.text)(\(parameters.map { "\($0.name): \($0.name)" }.joined(separator: ",")))
+                    let jsonValue = try JSONEncoder().encode(returnValue)
+                    return String(data: jsonValue, encoding: .utf8)!
+                }
+            }
+            """
+            tools.append((name: name.text, description: description,
+                          parameters: parameters,
+                          callableString: callableString,
+                          callableName: callableName.text))
+        }
+        
+        
+        return [
+            .init(extendedType: type,
+                inheritanceClause: .init(inheritedTypes: InheritedTypeListSyntax.init(arrayLiteral: .init(type: IdentifierTypeSyntax(name: "LlamaActor")))),
+                  memberBlock: """
+            {
+                \(raw: tools.map {
+                    $0.callableString
+                }.joined(separator: "\n"))
+            
+                static var tools: [String: (DynamicCallable, _JSONFunctionSchema)] {
+                    [\(raw: tools.map { tool in
+                        """
+                        "\(tool.name)": (\(tool.callableName)(), _JSONFunctionSchema(name: "\(tool.name)", description: "\(tool.description)", parameters: _JSONFunctionSchema.Parameters(properties: \(tool.parameters.count == 0 ? "[:]" : "[" + tool.parameters.map { parameter in
+                            """
+                            "\(parameter.name)": _JSONFunctionSchema.Property(type: \(parameter.type).self, description: "\(parameter.description)"),
+                            """
+                            }.joined() + "]"), required: [])))
+                        """
+                    }.joined(separator: ","))]
+                }
+            }
+            """)
+        ]
+    }
+}
+
+@main
+struct LlamaKitMacrosPlugin: CompilerPlugin {
+    let providingMacros: [Macro.Type] = [
+        LlamaActorMacro.self, ToolMacro.self
+    ]
+}

swift/main/main.swift

@@ -2,13 +2,6 @@ import LlamaKit
 import WeatherKit
 import CoreLocation
 
-@llamaActor actor MyLlama {
-    /// Get the current date.
-    public static func getCurrentDate() -> String {
-        Date.now.formatted(date: .long, time: .complete)
-    }
-}
-
 func downloadFile() async throws -> String {
     let fm = FileManager.default
     let tmpDir = fm.temporaryDirectory
@@ -36,6 +29,14 @@ func downloadFile() async throws -> String {
     return destinationURL.path()
 }
 
+
+@llamaActor actor MyLlama {
+    /// Get the current date.
+    @Tool public static func getCurrentDate() -> String {
+        Date.now.formatted(date: .long, time: .complete)
+    }
+}
+
 let params = GPTParams()
 params.modelPath = try await downloadFile()
 params.nPredict = 512

swift/test/LlamaKitTests.swift

@@ -78,24 +78,6 @@ struct LlamaGrammarSessionSuite {
 import WeatherKit
 import CoreLocation
 
-@llamaActor actor MyLlama {
-    struct CurrentWeather: Codable {
-        let temperature: Double
-        let condition: WeatherCondition
-    }
-    
-    /// Get the current weather in a given location.
-    /// - parameter location: The city and state, e.g. San Francisco, CA
-    /// - parameter unit: The unit of temperature
-    public static func getCurrentWeather(location: String, unit: String) async throws -> CurrentWeather {
-        let weather = try await WeatherService().weather(for: CLGeocoder().geocodeAddressString(location)[0].location!)
-        var temperature = weather.currentWeather.temperature
-        temperature.convert(to: .fahrenheit)
-        return CurrentWeather(temperature: temperature.value,
-                              condition: weather.currentWeather.condition)
-    }
-}
-
 func downloadFile() async throws -> String {
     let fm = FileManager.default
     let tmpDir = fm.temporaryDirectory
@@ -123,6 +105,25 @@ func downloadFile() async throws -> String {
     return destinationURL.path()
 }
 
+
+@llamaActor actor MyLlama {
+    struct CurrentWeather: Codable {
+        let temperature: Double
+        let condition: WeatherCondition
+    }
+    
+    /// Get the current weather in a given location.
+    /// - parameter location: The city and state, e.g. San Francisco, CA
+    /// - parameter unit: The unit of temperature
+    @Tool public static func getCurrentWeather(location: String, unit: String) async throws -> CurrentWeather {
+        let weather = try await WeatherService().weather(for: CLGeocoder().geocodeAddressString(location)[0].location!)
+        var temperature = weather.currentWeather.temperature
+        temperature.convert(to: .fahrenheit)
+        return CurrentWeather(temperature: temperature.value,
+                              condition: weather.currentWeather.condition)
+    }
+}
+
 @Test func llamaToolSession() async throws {
     let params = GPTParams()
     params.modelPath = try await downloadFile()