🧠

SwiftでRAG（Retrieval-Augmented Generation）システムを構築する

2024/11/05に公開

こんにちはStampの村本です。
最近はインバウンド向けの観光メディアtabisakiを開発したりしております。
ChatGPTに観光について聞くとそれなりに情報をもらえるんですが、案内をしてもらったりルートの提案ができるといいなと思いAIを開発しております。
今回は、すべてをSwiftで実装したRetrieval-Augmented Generation（RAG）システムについて紹介します。このシステムはデータ収集からベクトル検索、LLM（大規模言語モデル）の活用まで、Swiftだけで構成されています。RAGのワークフロー全体を通して、どのように各技術が組み合わさり、どのようにSwiftならではの利点を生かしているかを見ていきましょう。

 システム全体の構成RAGシステムは以下のようなプロセスで構成されています：
データの収集（クローラー）
収集したデータは、GitHubのSelenopsを使用して自動収集します。SelenopsはSwift製のクローラーで、シンプルなAPIで動作し、HTMLのデータ収集をスムーズに行えます。
データの保存（PostgreSQL, Fluent）
収集したデータは、Fluentを使ってPostgreSQLに保存します。FluentはSwiftで作られたORM（オブジェクト関係マッピング）で、データベースの操作を簡単に行えるため、Swiftエコシステムとの統合に最適です。
データの解析（LLM, Remark）
データ解析には、GitHubで提供されているOllamaKitとRemarkを使用します。これにより、収集したデータのテキスト解析とMarkdownへの変換が可能となります。RemarkはHTMLからMarkdownへの変換を行い、よりクリーンで扱いやすいデータ形式に整形します。
データのベクトル化（LLM）
テキストデータをベクトル化し、ベクトル検索を可能にするためにLLMを用います。ここでは主にユーザークエリに対する関連性の高い情報を効率的に検索するために使用されます。
データの取得（Vapor, LLM）
クライアントからのリクエストに応じて、Vaporを使用してデータを提供します。VaporはSwiftで実装されたWebフレームワークで、クライアントからのリクエストに対する迅速なレスポンスを実現します。

 利用したFramework
 VaporVaporは、Swiftで構築されたサーバーサイドWebフレームワークで、WebアプリケーションやAPIの開発を迅速かつ柔軟に行える環境を提供します。非同期処理や型安全性を活かし、パフォーマンスと信頼性の高いサーバーサイドアプリケーションを構築できます。

 RemarkRemarkは、HTMLコンテンツをMarkdownに変換するために設計されたSwiftライブラリです。Open Graph (OG) メタデータの抽出やフロントマターの自動生成もサポートし、静的サイトジェネレーターやMarkdownベースのアプリケーションに適した形式に整形します。

 OllamaKitOllamaKitは、SwiftでOllama APIとのやり取りを簡潔にするライブラリです。ネットワーク通信やデータ処理の複雑さをライブラリ内で処理し、Ollama APIの統合をシンプルかつ効率的に実現します。

 FluentFluentは、Swiftで書かれたORM（オブジェクト関係マッピング）ライブラリで、データベースとSwiftコードの橋渡しを行います。データベース操作をSwiftの型安全性と統合し、PostgreSQL、MySQL、SQLiteなど複数のデータベースをサポートしています。FluentはVaporフレームワークとシームレスに統合され、サーバーサイドSwiftアプリケーションでのデータ管理が容易になります。

 SelenopsSelenopsはSwift製の軽量なWebクローラーで、特定の単語をWebページ上で効率的に検索するために設計されています。Swift Concurrencyを活用しており、安全で高性能なクロールを実現しています。

 LAGのワークフローSwiftによるLAGの構成は非常にシンプルで、以下のような流れになります。
クローリングしたHTMLをMarkdownに変換
Selenopsで収集したHTMLデータをRemarkを使ってMarkdown形式に変換します。Markdownにすることで、テキストデータの処理やベクトル化が容易になります。
Markdownをベクトル化
Markdownに変換したデータをLLMを用いてベクトル化します。ベクトル化されたデータは、クエリと比較して関連性の高い情報を見つけ出すために使われます。
PostgreSQLに保存
ベクトル化されたデータはFluentを使ってPostgreSQLに保存します。これにより、保存と検索のパフォーマンスが最適化されます。
Fluentを通してクライアントから取得
クライアントからのリクエストに応じて、FluentとVaporを使ってデータを迅速に取得し、適切な情報を返します。

 クローラーの実装LAGシステムの最初のステップとなるクローラーについて説明します。クローラーは情報収集の要となる重要なコンポーネントです。

 データの流れまず、クローラーのデータの流れを理解しましょう：
URLから情報を収集
HTMLを解析してMarkdownに変換
データベースに保存
新しいURLを発見して次のクロール対象に追加
このシンプルな流れを実現するために、3つの主要なデータモデルを用意しています：
// Webページのリンク情報
final class Link: Model {
    @ID var id: UUID?
    @Field var url: String
    @Field var title: String
    @Field var crawledAt: Date
    @OptionalChild var pageContent: PageContent?
}

// ページの実際のコンテンツ
final class PageContent: Model {
    @ID var id: UUID?
    @Parent var link: Link
    @Field var title: String?
    @Field var body: String?
    @Field var crawledAt: Date
}

// クロール予定のURL
final class PagesToVisit: Model {
    @ID var id: UUID?
    @Field var url: String
    @Field var priority: Int
    @Field var retryCount: Int
}

 HTMLからMarkdownへの変換クローラーが収集したHTMLは、すぐにMarkdownに変換されます。これには自作のRemarkライブラリを使用しています：
let remark = try Remark(content)
let pageContent = PageContent(
    url: url.absoluteString,
    title: remark.title,
    body: remark.markdown,
    crawledAt: Date()
)
この変換により、後続の処理（ベクトル化やLLMでの分析）が容易になります。

 クローリングの制御クローラーの動作は、以下のようなシンプルなルールで制御されています：
HTTP/HTTPSのURLのみを対象とする
画像やPDFなどのバイナリファイルはスキップ
ニュースや記事など、優先度の高いコンテンツを先にクロール
エラー時は一定回数リトライ
クローラーのデリゲート
public protocol CrawlerDelegate: Actor {
    
    /// クローラーがURLを訪問するかを決定。
    func crawler(_ crawler: Crawler, shouldVisitUrl url: URL) -> Crawler.Decision
    
    /// クローラーがURLを訪問する前に通知。
    func crawler(_ crawler: Crawler, willVisitUrl url: URL)
    
    /// クローラーがすべての作業を完了したときに通知。
    func crawlerDidFinish(_ crawler: Crawler)
    
    /// クローラーに次の訪問URLを提供。
    func crawler(_ crawler: Crawler) async -> URL?
    
    /// クローラーが取得したHTMLコンテンツを通知。
    func crawler(_ crawler: Crawler, didFetchContent content: String, at url: URL) async
    
    /// クローラーがURLを訪問したことを記録。
    func crawler(_ crawler: Crawler, didVisit url: URL) async
    
    /// クロール中に発見したリンクを追加。
    func crawler(_ crawler: Crawler, didFindLinks links: Set<Crawler.Link>, at url: URL) async
    
    /// クローラーがURLをスキップしたことを通知。
    func crawler(_ crawler: Crawler, didSkip url: URL, reason: Crawler.SkipReason) async
}

 データの保存収集したデータはPostgreSQLに保存されますが、これはFluentを使って簡単に行えます：
try await database.transaction { db in
    let link = Link(url: url, title: title, crawledAt: Date())
    try await link.save(on: db)
    let content = PageContent(link: link, body: markdown)
    try await content.save(on: db)
}
このように、クローラーはシンプルながらも必要な機能を備えた設計となっています。次のステップであるベクトル化や検索のための基盤として、十分な機能を提供します。

 PostgreSQLのベクトル型対応RAGシステムを構築する上で大きな課題となったのが、PostgreSQLのベクトル型（pgvector）をSwiftから扱う方法でした。Fluentにはベクトルデータタイプのサポートがないため、独自の実装が必要でした。

 Vector型の実装まず、ベクトルの次元数を定義するためのプロトコルを用意します：
public protocol Dimension {
    static var count: Int { get }
}

public enum Dimensions {
    public enum D1536: Dimension { public static let count = 1536 }
    public enum D3072: Dimension { public static let count = 3072 }
}
この実装により、ベクトルの次元数を型レベルで保証できます。例えば、埋め込みモデルが1536次元のベクトルを生成する場合、Vector<D1536>として型安全に扱えます。
次に、PostgreSQLのvector型（OID: 23259）に対応するデータ型を定義します：
public struct Vector<D: Dimension>: PostgresCodable {
    public var value: [Float]
    
    public init(_ value: [Float]) throws {
        guard value.count == D.count else {
            throw VectorError.unexpectedDimension(
                expected: D.count,
                actual: value.count
            )
        }
        self.value = value
    }
}

 PostgreSQLとの連携PostgreSQLとの相互変換を可能にするために、PostgresCodableプロトコルを実装します：
extension Vector: PostgresCodable {
    public static var psqlType: PostgresDataType {
        .vector  // カスタム定義したベクトル型
    }
    
    // エンコード処理（Swiftの配列→PostgreSQLのベクトル）
    public func encode(
        into byteBuffer: inout ByteBuffer,
        context: PostgresEncodingContext<JSONEncoder>
    ) throws {
        let vectorString = "[\(value.map { String($0) }.joined(separator: ","))]"
        byteBuffer.writeString(vectorString)
    }
    
    // デコード処理（PostgreSQLのベクトル→Swiftの配列）
    public init<JSONDecoder: PostgresJSONDecoder>(
        from byteBuffer: inout ByteBuffer,
        context: PostgresDecodingContext<JSONDecoder>
    ) throws {
        // バイトバッファからFloatの配列を読み込む
    }
}

 実際の使用例このVector型は以下のように使用します：
struct Document: Model {
    @ID var id: UUID?
    @Field var title: String
    @Field var embedding: Vector<Dimensions.D1536>  // 1536次元のベクトル
}
このVector型の実装により、SwiftでもPostgreSQLのベクトル検索機能を活用できるようになりました。これはRAGシステムにおける重要な基盤となっています。

 Swiftでのベクトル検索の難しさSwiftには、ベクトル検索を処理するためのエコシステムが他言語に比べてまだ整っていません。例えば、PythonやGo、TypeScriptでは、LLMのワークフローを組むためのフレームワークとして、LangChainやGenkitなどが利用されています。しかし、Swiftにはこれらに対応するフレームワークが少なく、独自に実装する必要があります。

 オンデバイス環境でのAI
 オンデバイスで動作するLLMの現状MLX（Appleの機械学習フレームワーク）などを利用し、オンデバイスで動作するLLMモデルが増えてきていますが、まだ発展途上です。例えば、Llamaは比較的実行が容易なモデルですが、日本語対応や精度の面では課題が残ります。また、Apple IntelligenceのLLMも同様で、精度が十分でないとの報告が見られます。とはいえ、Appleが提供するSiri ShortcutのようなLLMワークフローは、Appleデバイス上でのシームレスな機能連携を可能にしており、これは他のAI企業には真似できない大きな強みです。

 AI搭載に向けた今後のビジネス戦略現在、SwiftでRAGを構築することは技術的に挑戦ですが、オンデバイスでのAI処理やLLMの進展は今後のビジネスに大きな影響を与えると考えています。AI技術の投資を続けることで、エッジデバイス上での効率的なデータ処理やリアルタイムの情報提供が可能となり、他社との差別化を図る大きな武器となります。
StampではAIに関する開発、ご相談などを受けております。

ご質問などあれば1amageekへご気軽にご連絡ください。

システム全体の構成

利用したFramework

Vapor

Remark

OllamaKit

Fluent

Selenops

LAGのワークフロー

クローラーの実装

データの流れ

HTMLからMarkdownへの変換

クローリングの制御

データの保存

PostgreSQLのベクトル型対応

Vector型の実装

PostgreSQLとの連携

実際の使用例

Swiftでのベクトル検索の難しさ

オンデバイス環境でのAI

オンデバイスで動作するLLMの現状

AI搭載に向けた今後のビジネス戦略

Discussion