Databricks 上でのナレッジグラフ RAG システムの構築、改善、デプロイ

Building, Improving, and Deploying Knowledge Graph RAG Systems on Databricks

公開日: 2025年4月1日

製造Less than a minute

によって Andrea Santurbano、チャンダナ・パドマナバン、Jiayi Wu 、ダン・ペチによる投稿

Summary

GraphRAGの概要: このブログ記事では、リトリーバル拡張生成（RAG）システムがNeo4jのようなグラフデータベースでどのように強化できるかを探ります。これにより、構造化データ内のエンティティ間の意味的な関係を捉えることで、より正確なAI出力を可能にします。
ユースケースとメリット: GraphRAGは、サイバーセキュリティにおける脅威検出、製造業における予知保全やサプライチェーン管理などの業界で応用でき、複雑なデータセットからより深い洞察を提供します。
Databricksでの実装: このブログ記事では、Neo4jを使用してDatabricks上にGraphRAGシステムを構築およびデプロイする方法の概要を説明します。LLM、Delta Tables、およびMosaic AI Agent Frameworkの統合を示し、エンドツーエンドのデプロイを実現します。

GraphRAGの理解

ナレッジグラフとは？

構造化データの別の表現方法ではなくナレッジグラフ（KG）を使用する理由を理解するには、ビジネス、人物、機械、顧客などのエンティティ間の明示的な関係と、それに関連付けられた属性または特徴に焦点を当てていることを認識することが重要です。高次元空間での類似性を優先する埋め込みやベクトル検索とは異なり、ナレッジグラフはデータポイント間の意味的な接続とコンテキストを表現することに優れています。ナレッジグラフの基本的な単位は事実です。事実は、次のいずれかの方法でトリプレットとして表すことができます。

HRT: <head, relation, tail>
SPO: <subject, predicate, object>

2つの簡単なKGの例を以下に示します。事実の左側の例は、<Andrea, loves, Irene>とすることができます。KGは、このような複数の事実のコレクションにすぎないことがわかります。しかし、グラフには意味があることに気づくかもしれませんが、左側の例は2人の間のロマンチックな関係を表していませんが、右側の例は2人の間のロマンチックな関係を表しています。

ナレッジグラフにおける意味の重要性を理解したところで、今後のコード例で使用するデータセットを紹介します。それはBloodHoundデータセットです。BloodHoundは、Active Directory環境内の関係と相互作用を分析するために設計された特殊なデータセットです。セキュリティ監査、攻撃パス分析、およびネットワーク構造における潜在的な脆弱性についての洞察を得るために広く使用されています。

BloodHoundデータセットのノードは、Active Directory環境内のエンティティを表します。これらは通常、以下を含みます。

ユーザー: ドメイン内の個々のユーザーアカウントを表します。
グループ: アクセス許可の割り当てのためにユーザーまたは他のグループをまとめるセキュリティグループまたは配布グループを表します。
コンピューター: ネットワーク内の個々のマシン（ワークステーションまたはサーバー）を表します。
ドメイン: ユーザー、コンピューター、およびグループを整理および管理するActive Directoryドメインを表します。
組織単位（OU）: ユーザーまたはグループなどのオブジェクトの構造化と管理に使用されるコンテナーを表します。
GPO（グループポリシーオブジェクト）: ドメイン内のユーザーおよびコンピューターに適用されるポリシーを表します。

ノードエンティティの詳細な説明はこちらでご覧いただけます。グラフ内の関係は、ノード間の相互作用、メンバーシップ、およびアクセス許可を定義します。エッジの詳細な説明はこちらでご覧いただけます。

The graph model of the BloodHound dataset

従来のRAGよりもGraphRAGを選択する場合

標準RAGよりもGraphRAGの主な利点は、取得ステップ中に正確な一致を実行できることです。これは、自然言語クエリの意味をダウンストリームのグラフクエリ言語で明示的に保持することによって可能になります。コサイン類似性に基づく密な検索技術は、クエリが完全に一致しない場合でも、あいまいな意味を捉え、関連情報を取得するのに優れていますが、精度が重要な場合があります。これにより、コンプライアンス、法律、または高度にキュレーションされたデータセットなど、あいまいさが許容されないドメインでGraphRAGは特に価値があります。

とはいえ、これら2つのアプローチは相互に排他的ではなく、それぞれの強みを活かすために組み合わされることがよくあります。密な検索は意味的な関連性のために広範囲をカバーできますが、ナレッジグラフは正確な一致または関係の推論によって結果を絞り込みます。

従来のRAGをGraphRAGよりも選択する場合

GraphRAGには独自の利点がありますが、課題もあります。主なハードルは問題を正しく定義することです。すべてのデータまたはユースケースがナレッジグラフに適しているわけではありません。タスクに高度に構造化されていないテキストが含まれる場合や、明示的な関係を必要としない場合は、追加の複雑さがそれに見合わない可能性があり、非効率的で最適ではない結果につながります。

もう1つの課題は、ナレッジグラフの構造化と維持です。効果的なスキーマを設計するには、詳細と複雑さのバランスをとるために慎重な計画が必要です。スキーマ設計が不十分だと、パフォーマンスとスケーラビリティに影響を与える可能性があります。また、継続的なメンテナンスにはリソースと専門知識が必要です。

リアルタイムパフォーマンスも別の制限事項です。Neo4jのようなグラフデータベースは、複雑なトラバーサルとマルチホップクエリのために、大規模または頻繁に更新されるデータセットでのリアルタイムクエリに苦労する可能性があり、密な検索システムよりも遅くなります。このような場合、ハイブリッドアプローチ（速度のために密な検索を使用し、クエリ後の分析のためにグラフを洗練する）が、より実用的なソリューションを提供できます。

GraphDBと埋め込み

Neo4jのようなグラフDBは、HNSWインデックスを介してベクトル検索機能を提供することもよくあります。ここでは、ベクトルデータベースよりも優れた結果を提供するために、このインデックスをどのように使用するかが異なります。クエリを実行すると、Neo4jはHNSWインデックスを使用して、コサイン類似性またはユークリッド距離などの尺度に基づいて最も近い一致する埋め込みを特定します。このステップは、ベクトル検索によって与えられる暗黙的な意味を活用して、クエリと意味的に一致するデータ内の開始点を見つけるために重要です。

グラフデータベースを際立たせているのは、この初期のベクトルベースの検索を強力なトラバーサル機能と組み合わせることができることです。HNSWインデックスを使用してエントリポイントを見つけた後、Neo4jはナレッジグラフの関係によって定義された明示的な意味を活用します。これらの関係により、データベースはグラフをトラバースして追加のコンテキストを収集し、ノード間の意味のある接続を明らかにすることができます。埋め込みからの暗黙的な意味とグラフの関係からの明示的な意味のこの組み合わせにより、グラフデータベースは、どちらかのアプローチだけでは達成できない、より正確でコンテキストが豊富な回答を提供できます。

DatabricksでのエンドツーエンドGraphRAG

GraphRAGは、複数のAIコンポーネントが連携して検索をよりスマートでコンテキストアウェアにするCompound AIシステムの実例です。このセクションでは、すべてがどのように連携するかを概観します。

GraphRAGアーキテクチャ

以下は、アナリストの自然言語の質問がNeo4jナレッジグラフから情報を取得する方法を示すアーキテクチャ図です。

GraphRAGを搭載した脅威検出のアーキテクチャは、DatabricksとNeo4jの強みを組み合わせています。

セキュリティオペレーションセンター（SOC）アナリストインターフェイス: アナリストはDatabricksを介してシステムと対話し、クエリを開始し、アラート推奨を受け取ります。
Databricks処理: Databricksはデータ処理、LLM統合を処理し、ソリューションの中心ハブとして機能します。
Neo4jナレッジグラフ: Neo4jはサイバーセキュリティナレッジグラフを保存および管理し、複雑な関係クエリを可能にします。

実装の概要

このブログでは、コードの詳細は省略します。完全な実装については、GitHubリポジトリを確認してください。GraphRAGエージェントを構築およびデプロイするための主要なステップを説明します。

Deltaテーブルからナレッジグラフを構築する: ノートブックでは、構造化データと非構造化データのシナリオについて説明しました。Neo4j Spark Connectorは、Unity Catalogのデータをグラフエンティティ（ノード/関係）に変換する非常に簡単な手段を提供します。
CypherクエリとQA用のLLMをデプロイする: GraphRAGは、クエリ生成と要約のためにLLMを必要とします。HuggingFaceのgpt-4o、llama-3.x、ファインチューニングされたtext2cypherモデルをデプロイし、プロビジョニングされたスループットエンドポイントを使用してそれらをサーブする方法を示しました。
GraphRAGチェーンを作成してテストする: GraphCypherQAChainを使用して、CypherおよびQA LLMとプロンプトに異なるLLMを使用する方法を示しました。これにより、MLflow Tracingを使用してグラスボックストレーシング結果でさらにチューニングできます。

Mosaic AI Agent Framework を使用してエージェントをデプロイ: Mosaic AI Agent Framework と MLflow を使用してエージェントをデプロイします。ノートブックでは、モデルのロギング、Unity Catalog への登録、サービングエンドポイントへのデプロイ、チャット用のレビューアプリの起動が含まれます。

結論

GraphRAG は、より決定的で文脈に関連性の高い AI 出力を生成するエージェントを構築するための強力かつ高度にカスタマイズ可能なアプローチです。ただし、その設計はケース固有であり、慎重なアーキテクチャと問題固有のチューニングが必要です。知識グラフを Databricks のスケーラブルなインフラストラクチャとツールと統合することで、構造化データと非構造化データをシームレスに組み合わせ、より深い文脈理解に基づいた実用的な洞察を生成するエンドツーエンドの Compound AI systems を構築できます。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事