Difyの外部ナレッジとしてDatabricks Mosaic AI Vector Searchを利用する

2025.02.07

Difyには外部のデータベースの結果をRAGに組み込める外部ナレッジの機能があります。今回は外部ナレッジとしてDatabricksのMosaic AI Vector Searchを利用できるようにします。

ライター：荒牧　大樹: 2007年ネットワンシステムズ入社し、コラボレーション・クラウド製品の担当を経て現在はAI・データ分析製品と技術の推進に従事。最近では次世代の計算環境であるGPU・FPGA・量子コンピュータに注目している。

Difyについて

Difyはノーコードで、RAGを使ったChatBotやワークフローを構築する事が出来ます。様々なシステムと連携する事が可能で、Local LLMやクラウドのLLMとの連携も可能です。今回はDifyの外部ナレッジ連携の機能を使ってDifyとDatabricksを連携させます。

Difyの外部ナレッジ機能について

Dify上で外部ナレッジを登録しておくと、入力されたPromptを元にQueryが作成されて、返ってきた文章を参考文献として文章生成を行う事が可能となります。外部ナレッジとして連携すると、Dify側からは以下の形式でQueryが外部ナレッジ側に送信されます。

POST <your-endpoint>/retrieval HTTP/1.1
-- ヘッダー
Content-Type: application/json
Authorization: Bearer your-api-key
-- データ
{
　"knowledge_id": "your-knowledge-id",
　"query": "your question",
　"retrieval_setting":{
　　"top_k": 2,
　　"score_threshold": 0.5
　}
}

外部ナレッジからは以下の形式でデータが返って来る事が期待されています。

HTTP/1.1 200
Content-type: application/json
{
　"records": [
　{
　　"metadata": {
　　　"path": "s3://dify/knowledge.txt",
　　　"description": "dify知識ドキュメント"
　　},
　"score": 0.98,
　"title": "knowledge.txt",
　"content": "これは外部知識のドキュメントです。"
　}
　]
}

DifyとDatabricksではデータの形が違う為、連携にはこれらの形式のJSONを受け取って、返信も決まった形式に変換するアプリが必要になります。

DifyとDatabricks Mosaic AI Vector Searchの接続

DifyとDatabricksを接続するアプリは様々な構成を取る事が出来ます。AWSのAPI GatewayとLambdaを利用するようなクラウド利用でも良いですし、Fast APIを使ってDifyをインストールしているサーバや、別サーバに独自に立てる事も可能です。今回はAPI GatewayとLambdaの構成としました。

変換アプリ作成時に以下の2点を注意してください。

Difyから飛んでくるQueryのURLには/retrievalが自動的に末尾に追加されます。
Difyから外部ナレッジ登録時に、Bodyが空のPOSTメッセージが飛んできます。アプリ側で空の場合もエラーを返さない処理が必要です。

以下Lambdaを利用する場合のサンプルコードです。API GatewayやToken認証の設定は適切に行ってください。Difyから受け取ったデータを元に、Databricks側を検索して結果を変換して返信しています。

import os
import json
from databricks.vector_search.client import VectorSearchClient

workspace_url = os.environ["WORKSPACE_URL"]
sp_id = os.environ["SP_ID"]
sp_secret = os.environ["SP_SECRET"]
vs_endpont_name = "vector_search_endpoint"
columns=["url","chunk","title"]

def find_relevant_doc(knowledge_id,query,retrieval_setting):
　# DatabricksのVector DBのサーバの設定を行います
　vsc = VectorSearchClient(service_principal_client_id=sp_id, service_principal_client_secret=sp_secret,workspace_url=workspace_url)　
　index = vsc.get_index(endpoint_name=vs_endpont_name, index_name=knowledge_id)

　# Databricks側で検索します。
　results = index.similarity_search(
　　　　　　query_text = query,
　　　　　　columns = columns,
　　　　　　num_results = retrieval_setting.get("top_k"),
　　　　　　score_threshold = retrieval_setting.get("score_threshold"),
　)
　return results

def get_records(results):
　#Databricksから得られた結果をDifyが望む形式に変換します。
　records = []
　for result in results["result"]["data_array"]:
　　record = {"metadata":{"path": result[0]},"content": result[1],"title": result[2],"score": result[3]}
　　records.append(record)
　return(records)

def lambda_handler(event, context):
　# API登録時の処理
　if not event.get("body"):
　　return {"statusCode": 200,"body": json.dumps({"message": "OK"})}
　try:
　　body = json.loads(event["body"])
　　knowledge_id = body.get("knowledge_id")
　　query = body.get("query")
　　retrieval_setting = body.get("retrieval_setting")

　　results = find_relevant_doc(knowledge_id,query,retrieval_setting)
　　records = get_records(results)
　
　　return{
　　　"statusCode": 200,
　　　"body": json.dumps({"records": records})
　　}

　except json.JSONDecodeError:
　　return{
　　　"statusCode": 400,
　　　"body": json.dumps({"error": "Invalid JSON"})
　　}

Dify上で外部ナレッジを登録する

作成したアプリは以下の手順でDifyに追加可能です。

ナレッジ>外部ナレッジAPI に登録
API EndpointにはAPI GatewayのURLを入力します。API KeyにはAPI Gatewayに設定したTokenを入力します。
ナレッジに登録
外部ナレッジ連携APIには先ほど作成した、外部ナレッジAPIを指定します。Knowledge_idにはDatabricks上でのVector DBを入力します。これは、先ほど作成したLambdaのコードでVector Searchの対象テーブルとしてKnowledge_idを利用するようにコードが記述されている為です。
検索テストを行う
検索画面からテストの検索を行います。結果の画面が出ています。