Vertex AI Searchのメタデータを利用してカテゴリごとにフィルタ検索をしてみる

2024.12.02 XIMIX 北脇

はじめに

本記事はNI+C TeamGCP Advent Calendar 2024の2日目の記事となります。

ドキュメント検索などのRAGを実装する場合、Vertex AI Agent Builder の Vertex AI Search機能（以下Vertex AI Searchという）を利用するパターンがあるかと思います。ドキュメントの数が数個のみであれば問題ないのですが、数が多くなってくると欲しいドキュメントが上位に表示されず、結果的によい回答結果が得られない、ということがあるかと思います。

今回ご紹介するメタデータを利用することで一つのアプリとデータストア内でドキュメントのカテゴリ分けを行い、実装の手間を削減しつつ、カテゴリを用いたフィルタリング機能を簡単に実行できるようになります。

メタデータを利用したカテゴリごとの検索

Vertex AI Search では、ドキュメントにメタデータを付与することで、柔軟な検索を実現できます。このメタデータを利用してドキュメントをカテゴリ分けし、フィルタ機能を組み合わせることで、特定のカテゴリのドキュメントのみを検索することが可能になります。

具体的には、ドキュメントをインデックスに追加する際に、metadata フィールドにカテゴリ情報を追加します。例えば、"technology"、"business"、"sports" などのカテゴリを設定できます。検索クエリを実行する際には、 filterパラメータを使用して、特定のカテゴリのドキュメントのみを絞り込むことができます。

この方法の利点は、一つのアプリとデータストアで複数のカテゴリを管理できるため、実装がシンプルになり、管理コストも削減できることです。また、ユーザーは簡単にカテゴリを切り替えて検索できるため、検索品質の向上にも繋がります。

実装例

それでは実装例をご紹介します。今回は例として「Cloud Run」と「Bigquery」の公式ドキュメントをPDFとしてダウンロードし、カテゴリ分けを行う例をご紹介します。

①Vertex AI Searchへの学習

下記ドキュメントをPDFでダウンロードしてGCS上に配置します。
https://cloud.google.com/bigquery?hl=ja
https://cloud.google.com/run/docs/overview/what-is-cloud-run?hl=ja

②メタデータの作成と学習

カテゴリごとにフィルタリングをするため、メタデータを作成します。メタデータの中身は下記の通りです。

{"id": "1", "structData": {"category": ["Bigquery"]}, "content": {"mimeType": "application/pdf", "uri": "gs://vertex-ai-search-metadata/BigQuery エンタープライズ向けデータウェアハウス _ Google Cloud.pdf"}} {"id": "2", "structData": {"category": ["CloudRun"]}, "content": {"mimeType": "application/pdf", "uri": "gs://vertex-ai-search-metadata/Cloud Run とは.pdf"}}

メタデータのパラメータは下記になります。

id:ドキュメントID
structData: 任意の情報
content: ドキュメント情報

この中でstructDataにcategoryを指定しており、ここで指定した値で出力結果のフィルタリングをすることが可能となっています。
作成したmetadataをGCS上にアップロードしVertex AI Searchに学習させます。

Vertex AI Searchのコンソール画面から新規にアプリ作成を行います。
手順は下記になります。

「アプリを作成する」押下

「Document Search」押下

「アプリ名」「会社名または組織名」を入力し「続行」押下

「データストアを作成」押下

「Cloud Storage」の「SELECT」押下

「リンクされた非構造化ドキュメント（JSONL止めたデータ）」を選択し、アップロードしたメタデータファイルを選択

データストア名に任意の名前を入力し、作成押下

作成したデータストアを選択し、「作成」押下

インポート実施後、実際に2件のPDFが学習されました。

②検索

では実際に検索してみましょう。まずはフィルタリングを行わない形で「cloud runとは？」と検索をして取得したファイル名のみを表示しています。
検索のサンプルソースは下記になります。

const { GoogleAuth } = require("google-auth-library");
const axios = require("axios");

async function getAccessToken() {
  const scope = "https://www.googleapis.com/auth/cloud-platform";
  const googleAuth = new GoogleAuth({
    scopes: scope,
  });

  const token = await googleAuth.getAccessToken();
  return token;
}

async function sendSearchByRestApiBySearchOnly(query, searchPath) {
  const url = `https://discoveryengine.googleapis.com/v1/${searchPath}`;
  const postData = {
    query: query,
    pageSize: 10,
    queryExpansionSpec: { condition: "AUTO" },
    spellCorrectionSpec: { mode: "AUTO" },
    contentSearchSpec: {
      snippetSpec: { returnSnippet: true },
      extractiveContentSpec: {
        maxExtractiveAnswerCount: 5,
      },
    },
  };
  const headers = {
    "Content-Type": "application/json",
    Authorization: "Bearer " + (await getAccessToken()),
  };

  try {
    const response = await axios.post(url, postData, {
      headers: headers,
    });
    for (const result of response["data"]["results"]) {
      if (
        result["document"]["derivedStructData"]["extractive_answers"] == null
      ) {
        continue;
      }

      const link = result["document"]["derivedStructData"]["link"];
      const linkSplit = link.split("/");
      const pdfName = linkSplit[linkSplit.length - 1];
      console.log("ファイル名: " + pdfName);
    }
  } catch (error) {
    console.log("error: " + error);
    throw error;
  }
}
let searchPath = `利用しているvertex ai searchのパスを指定`;
sendSearchByRestApiBySearchOnly("cloud runとは？", searchPath);

出力結果は下記になります。何もフィルタリングをしていないため2つのドキュメントが取得できました。
ファイル名: Cloud Run とは.pdf ファイル名: BigQuery エンタープライズ向けデータウェアハウス _ Google Cloud.pdf

次に、メタデータを利用したフィルタリングをしてみます。フィルタリングにはリクエストパラメータに「filter」のパラメータを追加します。
追記部分は赤字の個所となります。

const { GoogleAuth } = require("google-auth-library"); const axios = require("axios"); async function getAccessToken() { const scope = "https://www.googleapis.com/auth/cloud-platform"; const googleAuth = new GoogleAuth({ scopes: scope, }); const token = await googleAuth.getAccessToken(); return token; } async function sendSearchByRestApiBySearchOnly(query, filter, searchPath) { filter = 'category: ANY("' + filter + '")'; const url = `https://discoveryengine.googleapis.com/v1/${searchPath}`; const postData = { query: query, filter: filter, pageSize: 10, queryExpansionSpec: { condition: "AUTO" }, spellCorrectionSpec: { mode: "AUTO" }, contentSearchSpec: { snippetSpec: { returnSnippet: true }, extractiveContentSpec: { maxExtractiveAnswerCount: 5, }, }, }; const headers = { "Content-Type": "application/json", Authorization: "Bearer " + (await getAccessToken()), }; try { const response = await axios.post(url, postData, { headers: headers, }); for (const result of response["data"]["results"]) { if ( result["document"]["derivedStructData"]["extractive_answers"] == null ) { continue; } const link = result["document"]["derivedStructData"]["link"]; const linkSplit = link.split("/"); const pdfName = linkSplit[linkSplit.length - 1]; console.log("ファイル名: " + pdfName); } } catch (error) { console.log("error: " + error); throw error; } } let searchPath = `利用しているvertex ai searchのパスを指定`; sendSearchByRestApiBySearchOnly("cloud runとは？", "CloudRun", searchPath);
"CloudRun"というカテゴリでフィルタリングを行っています。フィルタ機能の使い方としては下記になります。
category: ANY("'フィルタしたいカテゴリ'")
上記のようにmetadataで指定したカテゴリ値を指定することによって検索対象のドキュメントを絞り込みすることが可能となります。出力結果は下記になります。
ファイル名: Cloud Run とは.pdf

BigQuery側のドキュメントは検索対象とはならず、CloudRun側のみのドキュメントを取得することができました。

まとめ

Vertex AI Search のフィルタ機能を利用することで、メタデータに基づいたカテゴリ検索を容易に実装できます。複数のアプリとデータストアを作成する必要がなく、一つの環境で複数のカテゴリを管理できるため、実装と管理の手間を大幅に削減できます。これにより、ユーザーはより効率的に必要な情報にアクセスできるようになり、検索エクスペリエンスの向上に貢献します。フィルタ機能を効果的に活用し、より高度な検索システムを構築してみてください。