Zenn
📌

Function CallingのためのSwift Macroを書いた

2024/11/22に公開
Swift
functioncalling
tech

前の記事

https://zenn.dev/fumitoito/articles/5e0dcd9883b0ff

最近のLLMサービスは関数定義を受け取り、必要に応じてクライアントに関数の呼び出しを依頼することで、より多様なタスクを実行できるようにする機能がある。呼び方は色々あるが、ChatGPTでは Function Calling、Anthropic Claudeでは Tool Use という名前で呼ばれている。

https://docs.anthropic.com/en/docs/build-with-claude/tool-use

Swiftで書かれたLLMクライアントでは、関数定義を表現する型をユーザーに記述させることでこの機能を実現している。例えばAnthropic ClaudeのAPIでは以下のようなJSONを渡すことで、APIから呼び出し可能な関数を公開することができる。

"tools": [
    {
        "name": "get_weather",
        "description": "Get the current weather in a given location",
        "input_schema": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The city and state, e.g. San Francisco, CA"
                }
            },
            "required": ["location"]
        }
    }
]

しかし、関数の定義を正確に表現するJSONを長期的にメンテナンスするのは現実的ではないだろう。関数のインターフェースは変更される可能性もあり、変更に追従していくのは簡単ではない。関数が多ければ「変更に追従しなくてはいけない」こと自体に気がつかない可能性もある。

既存のライブラリでは、FunctionCallingのためのstructやclassを提供することでこの問題に対応しようとしている。例えば、以下に示すのは MacPaw/OpenAPI でFunctionCallingを扱うコード例である。

https://github.com/MacPaw/OpenAI

// Declare functions which GPT-3 might decide to call.
let functions = [
  ChatFunctionDeclaration(
      name: "get_current_weather",
      description: "Get the current weather in a given location",
      parameters:
        JSONSchema(
          type: .object,
          properties: [
            "location": .init(type: .string, description: "The city and state, e.g. San Francisco, CA"),
            "unit": .init(type: .string, enumValues: ["celsius", "fahrenheit"])
          ],
          required: ["location"]
        )
  )
]

これはJSONがSwiftのオブジェクトに置き換わっただけで、関数定義の変更に追従するのが困難であるという本質的な問題は解決されていないように感じられた。

また、Function Callingが利用されるケースではAPIから呼び出し先の関数名と呼び出しパラメーターのセットが返されることになる。例としては以下のようなJSONが返される。

{
  "role": "assistant",
  "content": [
    {
      "type": "tool_use",
      "id": "toolu_01A09q90qw90lq917835lq9",
      "name": "get_weather",
      "input": {"location": "San Francisco, CA"}
    }
  ]
}

Swiftにはソースコードを生成する Swift Macros があるので、アノテーションをつけるだけでSwiftの関数から 関数定義を表現するJSON文字列が取得可能なプロパティ および APIからのレスポンスを受け取って実際の関数を呼び出すための関数 を自動生成できるのではないかと考えた。この思いつきを元に実装したのが FunctionCalling マクロである。

https://github.com/fumito-ito/FunctionCalling

この記事では FunctionCalling マクロについて簡単に紹介し、やや複雑なマクロを作成する上でどのような工夫をしたのかについて記載する。

FunctionCalling マクロについて

@FunctionCalling で修飾することで、Swiftのネイティブ関数をCodableなSwiftオブジェクトに変換したプロパティが実装される。また、LLMサービスからのレスポンスに対応してSwiftのネイティブ関数を呼び出す関数も同時に実装される。

// .claude, .chatGPT, .llamaOrGeminiをサポートしている
@FunctionCalling(service: .claude)
struct FunctionContainer {
    let session: URLSession
    let urlRequest: URLRequest

    @CallableFunction
    func getWeather(location: String) async throws -> String {
        let (data, _) = try await session.data(for: urlRequest)
        // TODO: data -> weather string
        return weather
    }
}

マクロによって以下のプロパティが自動的に実装され、APIに対するリクエスト/レスポンスを処理することができる。

extension FunctionContainer: ToolContainer {
    /// `@CallableFunction` で修飾した関数を表現するオブジェクトの配列
    /// `Encodable` を実装しているため、encodeしてJSONとして投げることができる
    var allTools: [Tool]

    /// `@FunctionCalling` で指定したサービスを表現する列挙体が返される
    /// 各列挙体は適切なencoder/decoderを返すので、 `allTools` のencode/decodeはそちらを使うのが無難
    var service: FunctionCallingService

    /// ほとんどのサービスでメソッド名と `[呼び出しパラメータ名: パラメータに渡す値]` のペアがAPIから返されるので、この関数に渡すことで適切な関数が呼び出される。
    func execute(method: String, parameters: [String: Any]) async throws -> String
}

マクロを書くときに工夫したこと

今回作成したマクロは複数の関数を取りまとめてコード生成をしたり、独自の型を受け取って独自の型を返すコードを生成したりするなど、一般的なexampleよりも複雑な要件を持っている。

また、内部的にもなるべく型安全に実装できるようにしたり、小さな単位でテストを書けるようにしたりなど、継続的にメンテナンスすることを念頭に幾つかの工夫がなされている。

本章ではその工夫についていくつか取り上げて紹介する。

適切な種類のマクロを選ぶ

Swiftのマクロは大きく Freestanding Macro と Attached Macro に分類され、Attached Macroに関しては実現するスコープに応じてさらに細かい分類がある。具体的には Swift Syntaxのドキュメント に記載がある。

  1. PeerMacro: 特定の型の隣接する位置に、補助的な要素（プロパティや関数など）を追加するためのマクロ。
  2. AccessorMacro: 特定のプロパティにゲッターやセッターを自動生成するためのマクロ。プロパティのアクセス制御やカスタム動作を設定する場合に利用する。
  3. MemberAttributeMacro: 型のメンバーに特定の属性を追加するマクロ。例えば、プロパティにデフォルトで付与されるべき属性を自動的に設定する際に使う。
  4. MemberMacro: 特定の型にメンバーを追加するマクロ。プロパティやメソッドを自動で生成して型に追加することができる。
  5. ExtensionMacro: 特定の型に対する拡張を自動生成するマクロ。既存の型に対して機能を追加したい場合に使用する。

当初、MemberMacroを利用して FunctionCalling マクロを実装する予定だったが、マクロを実装するインターフェースを確認すると、それが不可能であることがわかった。以下に示すのがMemberMacroが実装するprotocolである。

public protocol MemberMacro {
    static func expansion(
        of node: AttributeSyntax,
        providingMembersOf declaration: some DeclGroupSyntax,
        in context: some MacroExpansionContext
    ) throws -> [DeclSyntax]
}

実装から分かる通り、MemberMacroでは第2引数 declaration で渡される、マクロで修飾された対象（例えば関数）のシンタックスしか分からない。

今回はAPIからのレスポンスを受け取って実際の関数を呼び出す関数を実装する必要があるため、マクロスコープ外のどの関数がマクロで修飾されているのかを知る必要があるのだ。

struct FunctionContainer {
    // ↓このマクロからコード生成する場合には、
    // 自分自身以外にどの関数が `@CallableFunction` で修飾されているか分からない
    @CallableFunction
    func getWeather() -> String {
        // ...
    }

    @CallableFunction
    func getStockPrice() -> String {
        // ...
    }
}

MemberMacroではそのような要件は達成不可能なので、今回はExtensionMacroを採用した。Extension Macroであれば、型全体を拡張する形で各メンバー情報にアクセスできるため、今回の要件によりマッチしているからである。

また、Toolを宣言する型の中にも公開したい関数と公開したくない関数が混在できるようにした方が使いやすいだろうと考えた。そのため @FunctionCalling というExtension Macroと、 @CallableFunction というMember Macroの2つを定義した。@FunctionCalling のスコープに定義されている関数のうち、 @CallableFunction で修飾されている関数のみを対象として扱うようになっている。

struct FunctionContainer {
    @CallableFunction
    func getWeather() async throws -> String {
        // ...
    }

    // LLMには公開したくない共通処理などは `@CallableFunction` をつけなければ公開されない
    private func someFunction() async throws -> SomeResponse {
        // do something
    }
}

@CallableFunction は単なるアノテーションのためのマクロなので、実装的には何のコードも生成していない。

https://github.com/FunctionCalling/FunctionCalling/blob/099aa5df1eb0d3b0fb749f7a8c77d62c3b7a8708/Sources/FunctionCallingMacros/Macro/CallableFunctionMacro.swift#L14-L23

単なるアノテーションを実装するかは少し迷ったが、現実的なユースケースではLLMサービスに公開したいインターフェースと隠蔽したい実装が分離するであろうと考えてこのような構成になっている。

テストを書きやすい構成にする

基本的に、Swift Macros のテストは SwiftSyntaxMacroTestSupport を利用して書くことになる。
書き方の例としては以下のようなものだ。

assertMacroExpansion(
    """
    #stringify(a + b)
    """,
    expandedSource: """
    (a + b, "a + b")
    """,
    macros: testMacros
)

マクロの結果をテストするのであれば十分だが、実際にはマクロを展開するまでに多くの処理が必要で、その処理一つ一つをテストで確かめていきたいはずだ。
これを実現するには内部ロジックをユニットテスト可能な形に分離する必要がある。

Swiftのマクロはテストを含めたexampleが大量に公開されているものの、基本的にベタガキで実装されており、ある程度以上の規模のマクロを作成する際に「どのような構成でプロジェクトを構築するべきか」という情報がほとんど存在しなかった。

https://github.com/swiftlang/swift-syntax/tree/main/Examples/Sources/MacroExamples

しかし、Swiftのマクロで行うことは実にシンプルで、Syntaxを入力にして別のSyntaxを出力するだけだ。

Syntax -> （なんらかの処理） -> Syntax

そのため、マクロのエントリポイント, マクロ本体, Syntax Parser, Syntax Renderer, Entity のシンプルな構成を採用した。
こうすることで、Macro処理のどの時点で問題が発生しているのかが分かりやすくなり、各モジュール毎にテストを書くことで最低限の保守性を確保することができた。

具体的には FunctionCalling マクロでは以下のようにそれぞれの構成要素をtargetにしてSwift Packageで管理している。

targets: [
    // マクロのエントリポイント
    .target(name: "FunctionCalling", dependencies: [
        // 略
    ]),
    // マクロ本体
    .macro(name: "FunctionCallingMacros", dependencies: [
        // 略
    ]),
    // Renderer
    .target(name: "SyntaxRenderer", dependencies: [
        // 略
    ]),
    // Parser
    .target(name: "SyntaxParser", dependencies: [
        // 略
    ]),
    // Entity
    .target(name: "CommonModules"),
]

テストは各ターゲット毎に testTarget を指定できるので、マクロ本体に対しては SwiftSyntaxMacroTestSupport を利用し、その他のターゲットに対しては XCTest でテストを記述している。

SyntaxParserのテストを記述する場合には、 swift-syntax に含まれる SwiftParser を利用することで文字列を直接Syntaxに変換できる。これで比較的簡単にテストを書くことができた。

import SwiftParser

class SampleTest: XCTestCase {
    func testIsCallableFunction() throws {
        let classString = """
        @FunctionCalling(service: .claude)
        class Sample {
            @CallableFunction
            func concat(words: [String]) -> String {
                return ""
            }
        }
        """

        let syntax = Parser.parse(source: classString)

        // 自分で定義したParserがSwiftSyntaxを受け取って適切なEntityを生成できるかをテストする
        let declarations = MyCallableFunctionParser.parse(syntax: syntax)
        let contat = try XCTUnwrap(declarations.first)

        XCTAssertEqual(concat.name, "concat")
        XCTAssertEqual(concat.description, "")
        XCTAssertEqual(concat.parameters.count, 1)
    }
}

SyntaxRendererはEntityを受け取って文字列を出力するような構成にしておくとテストが書きやすかった。

func testFillingTemplateWithSingleFunction() throws {
    let model = Execute(functions: [Self.getHTML])

    let filledTemplate = try ExecuteSyntax.render(with: model)
    let expect = """
    func execute(methodName: String, parameters: [String: Any]) async throws -> String {
        do {
            switch methodName {
            case "getHTML":
                let urlString = parameters["urlString"] as! String
                return try await Self.getHTML(
                    urlString: urlString
                ).description
            default:
                throw FunctionCallingError.unknownFunctionCalled
            }
        } catch let error {
            return error.localizedDescription
        }
    }
    """
    XCTAssertEqual(filledTemplate, expect)
}

マクロ全体を型安全に実装する

前項で記述した通り、Swiftのマクロを実装する場合には マクロ本体 ではなく マクロのエントリポイント がライブラリとして公開される。Swift Packageの記述としては以下のようになる。

products: [
    // ライブラリとして公開されるのはマクロのエントリポイント
    .library(name: "FunctionCalling", targets: ["FunctionCalling"])
],
targets: [
    // マクロのエントリポイント
    // このターゲットはマクロ本体に依存する
    .target(name: "FunctionCalling", dependencies: [
        "FunctionCallingMacros",
        // 略
    ]),
    // マクロ本体
    // マクロ本体はSyntaxParserやSyntaxRenderer, Entityに依存する
    .macro(name: "FunctionCallingMacros", dependencies: [
        "SyntaxParser",
        "SyntaxRenderer",
        "CommonModules"
    ]),
    // 略
]

各モジュールが共通のEntityを参照しているので型安全に実装できそうに見えるが、FunctionCalling のようなマクロでは1点問題があった。

以下は @FunctionCalling で修飾したclassなどに実装されるコードだが、この中で利用されている型やProtocolはライブラリ内部で利用するとともに外部にも公開したいEntityである。

/// `ToolContainer` や `Tool`, `FunctionCallingService` はSyntaxParserやSyntaxRendererで利用したいし、ライブラリのユーザーにも公開したい
extension FunctionContainer: ToolContainer {
    var allTools: [Tool]
    var service: FunctionCallingService
    func execute(method: String, parameters: [String: Any]) async throws -> String
}

マクロ内の処理を型安全に実装するにはこれらのEntityはSyntaxParserやSyntaxRendererから参照できる必要がある。外部からこれらのEntityを利用できるようにするために、Entityが定義されている CommonModules を公開しても良いが、そうすると公開したくないEntityまで公開されてしまうというジレンマがあった。

同じように独自の型を引数として受け取るマクロである swift-testing のコードを参照すると、概ね全ての処理をマクロのエントリポイント側に実装し、マクロ本体はマクロのエントリポイントで実装している関数を呼び出すコードを文字列で生成するというストロングスタイルを採用していた。

例としては以下のように Testing.Tag.__fromStaticMember(of:) が呼び出し側から見られる状態にあるという前提に基づいてコード生成をしている。

https://github.com/swiftlang/swift-testing/blob/93562395a54ed0c019355a235136570b9ca92451/Sources/TestingMacros/TagMacro.swift#L84-L90

これ自体が間違っているとは当然思わないが、なるべく型安全にコードを書いていきたいという自分の目標とは異なる思想だった。

結局、正解と思える解決策を考えつくことができず、仕方なくtypealiasで誤魔化すというワークアラウンドをとっている。

https://github.com/FunctionCalling/FunctionCalling/blob/099aa5df1eb0d3b0fb749f7a8c77d62c3b7a8708/Sources/FunctionCalling/PublicModules.swift#L9-L16

頑張ってSwiftSyntaxを理解しようとしない

SyntaxRendererを実装しようとする際にSwiftSyntaxの馴染みなく煩雑なAPIを理解するのは大きな苦痛を伴う作業である。

// extensionでEquatableを実装したいだけなのに、この文法を理解する必要があるのは厳しい
ExtensionDeclSyntax(
    extensionKeyWord: TokenSyntax(
        .identifier("Equatable")
    ),
    extendedType: type
) {}

単にマクロが実装したい場合であれば SwiftSyntaxMacros をインポートするだけで苦痛を逃れることができる。各DeclSyntaxに文字列を渡すだけのコンストラクタが追加されるのだ。

import SwiftSyntaxMacros

// こういうのでいいんだよ
try ExtensionDeclSyntax("extension \(type.trimmed): Equatable {}")

もちろんSwiftSytaxの生のAPIを触ったほうがより型安全に実装することが可能だし、ループや分岐などを扱いやすくなる。

今回、筆者は複雑なコード生成のためにループや分岐などのロジックが必要だったが、SwiftSytanxのAPIを覚えるのは不可能であると早々に諦めてしまったため、テンプレートエンジンを採用することで問題に対処した。

https://github.com/groue/GRMustache.swift

これは筆者がメンテナを勤めているテンプレートエンジンで、ロジックがシンプルで手に馴染んでいるので利用した。以下は簡略化した例だが、ifやforのような簡単なロジックを扱えるテンプレートにオブジェクトを渡して展開している。

enum ExecuteSyntax {
    /// Entityをテンプレートに対して展開してSyntaxとして返したい文字列を生成している
    static func render(with templateObject: Execute) throws -> String {
        let template = try Template(string: ExecuteTemplate.templateString)
        return try template.render(templateObject).emptyLineRemoved
    }
}

let templateString: String =
"""
func execute(methodName: String, parameters: [String: Any]) async -> String {
    do {
        switch methodName {
{{#functions}}
        case "{{method_name}}":
{{#parameters}}
{{#required}}
            let {{name}} = parameters["{{name}}"] as! {{type}}{{! Strongly typed cast }}
{{/required}}
{{^required}}
            let {{name}} = parameters["{{name}}"] as? {{type}}{{! Optional cast without default }}
{{/required}}
{{/parameters}}
{{#is_static}}
            return try await Self.{{method_name}}(
{{/is_static}}
{{^is_static}}
            return try await {{method_name}}(
{{/is_static}}
{{#parameters}}
{{#isOmitting}}
                {{name}}{{^isLast}},{{/isLast}}
{{/isOmitting}}
{{^isOmitting}}
                {{name}}: {{name}}{{^isLast}},{{/isLast}}
{{/isOmitting}}
{{/parameters}}
            ).description
{{/functions}}
        default:
            throw FunctionCallingError.unknownFunctionCalled{{! Unknown function error }}
        }
    } catch let error {
        return error.localizedDescription
    }
}
"""

最終的にはSwiftSyntaxのAPIに対応し、脱テンプレートエンジンすることを目指しているが、さっさとマクロを実装するには必要十分であったし、テンプレートエンジン側をSwift6へ対応するモチベーションにも繋がったので結果としては良かった。

SwiftSyntaxのExtensionを用意する

最後もSwiftSyntaxの愚痴になるが、とにかくAPIがプリミティブすぎる。たとえば DeclGroupSyntax にぶら下がる関数シンタックス FunctionDeclSyntax のリストを取り出す処理でさえ、自分で書かなくてはいけないのだ。

// DeclGroupSyntaxにぶら下がるFunctionDeclSyntaxくらいは
// プリミティブなAPIとして実装してほしいが、そんな贅沢なものはない
let functions = DeclGroupSyntax().functions // <= エラー

// こんな感じのExtensionを書いていくことになる
extension DeclGroupSyntax {
    var functions: [FunctionDeclSyntax] {
        return memberBlock.members.compactMap({ $0.decl.as(FunctionDeclSyntax.self) })
    }
}

この手の問題はマクロを書こうとした多くの人間が直面しているようで、OSSにはいくつものSwiftSyntax Extensionが公開されている。

筆者は主に https://github.com/IanKeen/MacroKit を参考にしつつ自前のExntensionを書いた。これはどのようなマクロを書く場合でも直面する問題なので、これからマクロを書こうと考えている方はOSSの実装にサッと目を通しておくと問題を回避しやすいだろう。

まとめ

この記事では @FunctionCalling マクロで実現したこと、Swift Macroを実装するにあたって工夫した点を記載した。

適切なマクロの種類を選択し、テストを書きやすい構造を保ち、型安全に実装できるようにすることで長期的にメンテナンス可能なマクロを実装できた。また、マクロを実装するにあたってSwiftSyntaxに対する深い理解は必ずしも必要ではないと実感できた。

おまけ

@FunctionCalling を実装するにあたり、関数に記載されたSwift DocのSyntaxを構造化されたSwiftオブジェクトに変換するライブラリを書いた。

https://github.com/fumito-ito/DocumentationComment

これはFunction Callingを行う際に description に関数や引数の説明を記載することでLLMが関数を利用する精度を高める機能を利用するためである。
Swift MacroやSwift Syntaxはこういった「普段は気にも留めないが、実は役に立つかもしれない」ものに気がつかせてくれるテクノロジーであると感じた。

これからも機会があればOSSとして公開していきたい。

Discussion