RAG Engine クロス コーパス検索

重要: RAG コーパスを作成するときは、コーパスの description を設定します。作成後に description を編集することはできません。質の高い説明は、コーパス間の効果的な検索に不可欠 です。システムは `description` に基づいてクエリに適したコーパスを選択するためです。

このページでは、バックエンドで Agentic Retrieval を使用して RAG コーパス間検索をサポートする 2 つの API（AsyncRetrieveContexts と AskContexts）について説明します。

AsyncRetrieveContexts API: 複数の RAG 管理コーパスから関連するコンテキストを取得できる非同期 API です。これは長時間実行 API です。ユーザーはオペレーション ID を使用して、クエリのステータスと結果を取得できます。

AskContexts API: 複数の RAG 管理コーパスを検索して、クエリに対する回答を直接生成する同期 API です。

重要: この機能を使用するには、プロジェクトの [Vertex AI ユーザー] ロールを RAG Engine サービス アカウントに付与する必要があります: service-[Your project's automatically generated project number]@gcp-sa-vertex-rag.iam.gserviceaccount.com。

重要: この機能は us-central1 でのみ使用できます。

システム アーキテクチャとコーパスの選択

このセクションでは、RAG コーパス間検索システムのアーキテクチャと、検索に適したコーパスを選択する方法について説明します。

システム アーキテクチャ

コーパス間検索システムは、エージェント アプローチを使用して、複数のコーパスにわたる検索を調整します。アーキテクチャの概要には、次のコンポーネントが含まれます。

  • オーケストレータ/ルーター: ユーザー クエリを受け取り、検索プロセスを調整します。
  • プランニング エージェント: クエリを分析し、descriptions に基づいて最も関連性の高いコーパスを特定します。
  • 検索エンジン: 選択した RAG コーパスに対してセマンティック検索を実行します。
  • 推論エージェント: 取得したコンテキストがクエリに回答するのに十分かどうかを評価します。不十分な場合は、十分なコンテキスト認識（SCA）に基づいて別の検索ループをトリガーするフィードバックを生成します。詳細については、SCA のブログ投稿と研究論文をご覧ください。
  • LLM ジェネレータ（AskContexts の場合）: 取得したコンテキストを合成して最終的な回答を生成します。
Cross Corpus Retrieval アーキテクチャ

コーパスの選択方法

複数の RAG 管理コーパスに対してクエリが実行される場合、システムはクエリに回答するために必要な情報を含むコーパスを特定する必要があります。これは、Agentic Retrieval によって実現されます。

  1. コーパス マッピング: システムは、利用可能な RAG コーパスのマップと、その技術的な説明（RAG コーパスの description フィールドで指定）を保持します。
  2. セマンティック マッチング: 戦略プランニング エージェント（Gemini モデルを搭載）は、利用可能なコーパスの descriptions に対してユーザーのクエリを評価します。
  3. ターゲット ルーティング: プランナーは、クエリの特定の部分を最も関連性の高いコーパスにマッピングする検索プランを生成します。これにより、すべてのコーパスを単純に検索するのではなく、検索を絞り込み、効率化できます。

コードサンプル

次のコードサンプルは、RAG コーパス間検索 API の使用方法を示しています。

AsyncRetrieveContexts API

このコードサンプルは、複数の RAG コーパスからコンテキストを非同期で取得する方法を示しています。

Python 

from vertexai.preview import rag
import vertexai

PROJECT_ID = "PROJECT_ID"
LOCATION = "LOCATION"
RAG_CORPUS_1_ID = "RAG_CORPUS_1_ID"
RAG_CORPUS_2_ID = "RAG_CORPUS_2_ID"

# Initialize API once per session
vertexai.init(project=PROJECT_ID, location=LOCATION)

# Async retrieve contexts from multiple corpora
# This can be run in an environment like a Colab notebook.
response = await rag.async_retrieve_contexts(
    text="Why is the sky blue?",
    rag_resources=[
        rag.RagResource(
            rag_corpus=f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{RAG_CORPUS_1_ID}",
        ),
        rag.RagResource(
            rag_corpus=f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{RAG_CORPUS_2_ID}",
        ),
    ],
)
print(response)

REST 

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RAG_CORPUS_1_ID="RAG_CORPUS_1_ID"
RAG_CORPUS_2_ID="RAG_CORPUS_2_ID"
QUERY="Why is the sky blue?"

curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  "https://${LOCATION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}:asyncRetrieveContexts" \
  -d '{
    "query": {
      "text": "'"$QUERY"'"
    },
    "tools": {
      "retrieval": {
        "disable_attribution": false,
        "vertex_rag_store": {
          "rag_resources": {
            "rag_corpus": "projects/'"${PROJECT_ID}"'/locations/'"${LOCATION}"'/ragCorpora/'"${RAG_CORPUS_1_ID}"'"
          },
          "rag_resources": {
            "rag_corpus": "projects/'"${PROJECT_ID}"'/locations/'"${LOCATION}"'/ragCorpora/'"${RAG_CORPUS_2_ID}"'"
          }
        }
      }
    }
  }'

オペレーションの結果を取得する

AsyncRetrieveContexts は長時間実行オペレーションであるため、GetOperations API を使用して、オペレーション ID を使用して結果をポーリングできます。

Python 

from google.cloud import aiplatform_v1beta1

PROJECT_ID = "PROJECT_ID"
LOCATION = "LOCATION"
OPERATION_ID = "OPERATION_ID"

client = aiplatform_v1beta1.VertexRagServiceClient(
    client_options={"api_endpoint": f"{LOCATION}-aiplatform.googleapis.com"}
)

operation = client.get_operation(
    request={"name": f"projects/{PROJECT_ID}/locations/{LOCATION}/operations/{OPERATION_ID}"}
)

if operation.done:
    print(operation.response)
else:
    print("Operation is still running")

REST 

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
OPERATION_ID="OPERATION_ID"

curl -X GET \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  "https://${LOCATION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/operations/${OPERATION_ID}"

AskContexts API

このコードサンプルは、複数の RAG コーパスを同期的に検索して回答を直接生成する方法を示しています。

Python 

from vertexai.preview import rag
import vertexai

PROJECT_ID = "PROJECT_ID"
LOCATION = "LOCATION"
RAG_CORPUS_1_ID = "RAG_CORPUS_1_ID"
RAG_CORPUS_2_ID = "RAG_CORPUS_2_ID"

# Initialize API once per session
vertexai.init(project=PROJECT_ID, location=LOCATION)

# Ask contexts from multiple corpora
response = rag.ask_contexts(
    text="Why is the sky blue?",
    rag_resources=[
        rag.RagResource(
            rag_corpus=f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{RAG_CORPUS_1_ID}",
        ),
        rag.RagResource(
            rag_corpus=f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{RAG_CORPUS_2_ID}",
        ),
    ],
)
print(response)

REST

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RAG_CORPUS_1_ID="RAG_CORPUS_1_ID"
RAG_CORPUS_2_ID="RAG_CORPUS_2_ID"
QUERY="Why is the sky blue?"

curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  "https://${LOCATION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}:askContexts" \
  -d '{
    "query": {
      "text": "'"$QUERY"'"
    },
    "tools": {
      "retrieval": {
        "disable_attribution": false,
        "vertex_rag_store": {
          "rag_resources": {
            "rag_corpus": "projects/'"${PROJECT_ID}"'/locations/'"${LOCATION}"'/ragCorpora/'"${RAG_CORPUS_1_ID}"'"
          },
          "rag_resources": {
            "rag_corpus": "projects/'"${PROJECT_ID}"'/locations/'"${LOCATION}"'/ragCorpora/'"${RAG_CORPUS_2_ID}"'"
          }
        }
      }
    }
  }'

次のステップ