🌊

LangChain on Vertex AI(プレビュー) で Vertex AI Search と RAG する

2024/07/10に公開

LangChain のマネージドサービスの発表

Google Cloud Next'24 Las Vegas で LangChain on Vertex AI(プレビュー) が発表されました。
LangChain on Vertex AI は Reasoning Engine と呼ばれるマネージドサービスを利用して、LangChain を利用した AI エージェントを効率よく開発、運用できることを目指しています。
なお Reasoning Engine 自体は、LangChain はもちろん他のフレームワークや DIY した Python アプリケーションもサポートしています。

何をしてくれるのか

これまで LangChain アプリケーションを例えば Cloud Run で実装する場合、以下のような知識が必要となっていました。

  • コンテナと Cloud Run(インフラ)、セキュリティやスケールなどの運用
  • FastAPI や Django などによる Serving や、そのチューニング
  • 開発、本番環境でのログ、メトリクスなどの運用

これらをできるだけ不要にして、開発に集中できる環境を実現します。
LangChain on Vertex AI が提供するものは以下のとおりです。

  • LangChain によるエージェント開発を楽にするテンプレート
  • コードを使った本番環境へのかんたんなデプロイ
  • マネージドな API Serving
  • 開発、本番環境の一貫した運用(ログ、メトリクス、トレース)

特に Gemini の Function Calling が LangChain だけを使うよりかんたんに実装できるのは魅力的です。例えば Function Calling を使って、Python から接続できるデータストアを RAG することもシンプルに実現できます。また自分でインフラを用意することなくマネージドな API サーバーで運用してくれるので、インフラを考えたくない LangChain / Python プログラマーの方におすすめできます。

制約事項

LangChain on Vertex AI(及び Reasoning Engine) 自体はプレビュー(2024年7月時点)となっており、制約事項があります。

  • 利用可能リージョンは us-central1 のみ
  • Python のみをサポート(Python 3.8〜3.11 が必要)
  • Gogole Cloud コンソールや、gcloud などから操作ができず、SDKを使ったコードや REST APIからオペレーションをする必要がある

またひとつのクエリーで2つ以上の Function Calling が発生する場合も 400 エラーとなってしまう事象を確認しています。ご注意ください。

今回試すこと

今回は LangChain で動作確認のためのデモアプリを作りながら、Gemini の Function Calling を通じて Vertex AI Search を RAG して連携するアプリを実装してみます。

Vertex AI Search の準備

Vertex AI Search の準備はこちらを参照してご準備ください。
今回は BigQuery に Gemini で生成した架空のゲームの武器の情報を詰め込んでおり、インターネットや現実世界にはそのままはないであろうデータを用意しています。これはクエリーへの回答が Vertex AI Search から引き出していることを確認しやすくするためです(もちろん私の趣味もあります :)。

開発開始

1. Function Calling で利用する関数を準備

Function Calling に利用する関数は、ふつうの Python の関数で OK です。
関数には Gemini が文脈に応じて使用するツールを選定できるようにするため、関数定義内のコメントで、この関数の役割やシグネチャを説明します。
LangChain で同等のことをする場合に必要な description や デコレータは必要ありません。
関数なのでほぼなんでもできますが、例えばシンプルに外部の API を呼び出したい場合にも API コールをする関数を定義することで実現します。
ここでは関数の役割だけを記述していますが、引数や戻り値の情報もいれるとより意図したとおりに関数を Call させることができるでしょう。

main.py
def search_sword_items(query: str) -> str:
    """
    Search for information about a specific sword / katana
    """
    from langchain_google_community import VertexAISearchRetriever
    print(f"called search_sword_items with arg {query}") # 呼ばれてるのを確認するため
    retriever = VertexAISearchRetriever(
        project_id=PROJECT_ID, # プロジェクトID
        data_store_id=DATASTORE_ID, # データストアのID
        location_id=DATASTORE_LOCATION, # データストアの場所(global)
        engine_data_type=1, # データの種類は1(構造データ)を指定
        max_documents=10, # 最大ドキュメント数
        get_extractive_answers=False, # Extractive Answer は無効
    )

    result = str(retriever.invoke(query))
    return result

Function Calling が上手く動いているか確認するため、もうひとつ関数を用意してみます。
[こちら](https://openweathermap.org/api)のサービスを利用して、指定された地域の現在の天気予報情報を外部の API から取得します。

main.py
def get_current_weather(location: str) -> dict
    """
    Get the current weather information for a specific location
    Args:
        location must be specified in lower-case English
    """
    
    import requests
    print(f"called get_current_weather with arg {location}") # 呼ばれてるのを確認するため
    response = requests.get(
        f"http://api.weatherapi.com/v1/current.json?key={WEATHERAPI_APIKEY}&q={location}&aqi=no"
    )
    return response.json()

ひとつは仮想の世界の情報をもったデータストアを検索する関数、ひとつは現実世界の今を取得する関数なので、なんの用途のアプリケーションかよくわからない例になりましたが、ご容赦ください。

2. LangChain テンプレートを使って、アプリケーションを開発

Reasoning Engine で提供されている LangChainAgent クラスを利用します。

main.py

import vertexai
from vertexai.preview import reasoning_engines

vertexai.init( # 初期化
    project=PROJECT_ID, # プロジェクトID
    location="us-central1", # リージョン(現在は us-central1 のみサポート)
    staging_bucket=f"gs://{BUCKET_NAME}", # デプロイの際に利用される Cloud Storage のバケットを指定(ない場合は作成してください)
)

model = "gemini-1.5-flash" # モデルは Gemini 1.5 Flash を指定

model_kwargs = {
    "temperature": 0.1,
    "safety_settings": safety_settings, # 安全性設定は別途定義していますが、長いので省略
}

agent = reasoning_engines.LangchainAgent(
    model=model,
    tools=[ # 利用したい関数を配列に登録
        get_current_weather,
        search_sword_items,
    ],
    model_kwargs=model_kwargs,
)

3. アプリケーションをローカルでテスト

まず、ローカルで SDK を利用しているコードに権限を与えます。

gcloud auth application login

これでブラウザで Google Cloud にサインインしている権限を付与することができました。

先ほどのコードに続いて、テストするコードを書いて試してみます。

main.py
response = agent.query(input="鬼切丸と氷の刃 フロストはどっちが強いですか?")

出力例

鬼切丸は攻撃力が135.0でレアリティはA+、氷の刃フロストは攻撃力が110.0でレアリティはBです。鬼切丸の方が攻撃力が高くレアリティも高いので、鬼切丸の方が強いと言えるでしょう。

Vertex AI Search に登録した情報を参照した上で、Gemini による回答が生成できているようです。
ちなみにこのときに関数にわたっている引数は、"鬼切丸と氷の刃 フロスト 比較" となっていました。

文脈に応じて別の関数も呼んでくれるか確認します。

main.py
response = agent.query(input="京都の天気を教えて")

出力例

京都の天気は、曇り時々晴れで、気温は29.2度です。体感温度は35.3度です。風速は6.1km/hで、
南西の風です。湿度70%、降水量は0.0mmです。

なお利用している Weather API の地域名は英単語で渡す必要がありますが、関数のコメントで日本語で与えたものを英語の小文字にするよう指示しているので、日本語で指示しても変換して渡してくれています。関数にわたっている引数は、"kyoto"となっていました。

振り分けできています。とてもかんたんですね。

4. アプリケーションを Google Cloud にデプロイ

ローカルでテストできたので、Google Cloud の Vertex AI にデプロイします。

まず、Reasoning Engine から Vertex AI Search にクエリーをするので、Vertex AI Search が利用するサービスアカウントに事前に IAM 権限を与えておきます。

PROJECT_ID=<プロジェクトID>
PROJECT_NUMBER=<プロジェクト番号>

gcloud projects add-iam-policy-binding ${PROJECT_ID} \
    --member=serviceAccount:service-${PROJECT_NUMBER}@gcp-sa-aiplatform-re.iam.gserviceaccount.com \
    --role=roles/discoveryengine.editor

デプロイします。

main.py
DISPLAY_NAME = "My Sample App"

reasoning_engines.ReasoningEngine.create(
    agent, # 上記で作成した agent
    requirements=[
        "google-cloud-aiplatform[reasoningengine,langchain]",
        "langchain",
        "langchain-core",
        "google-cloud-discoveryengine",
        "requests==2.*",
        "langchain-google-community",
        "langchain-google-vertexai",
    ],
    display_name=DISPLAY_NAME,
)

出力例

Using bucket shingo-ar-dev0622
Writing to gs://shingo-ar-dev0622/reasoning_engine/reasoning_engine.pkl
Writing to gs://shingo-ar-dev0622/reasoning_engine/requirements.txt
Creating in-memory tarfile of extra_packages
Writing to gs://shingo-ar-dev0622/reasoning_engine/dependencies.tar.gz
Creating ReasoningEngine
Create ReasoningEngine backing LRO: projects/11092305173/locations/us-central1/reasoningEngines/7418448933069783040/operations/2103212671301058560
ReasoningEngine created. Resource name: projects/11092305173/locations/us-central1/reasoningEngines/7418448933069783040
To use this ReasoningEngine in another session:
reasoning_engine = vertexai.preview.reasoning_engines.ReasoningEngine('projects/11092305173/locations/us-central1/reasoningEngines/7418448933069783040')

デプロイには3分程度かかります。
出力された Reasoning Engine の Resource name をメモしておきましょう。

ログは自動で Cloud Logging に転送されています。
ちなみにログを見ると、Python のアプリケーションサーバーとして uvicorn が起動していることがわかると思います。

5. テスト

では、Google Cloud にデプロイ後のアプリケーションをテストしてみます。
新しくクライアント用の Python ファイルを準備して、ローカルの環境から実行します。実行には Vertex AI Reasoning Engine サービスエージェントの権限が必要になりますが、冒頭で gcloud でサインインしているので権限は含まれているものとして、省略します。

id には先ほどメモしておいた、Reasoning Engine の Resource name を指定しています。

client.py
from vertexai.preview import reasoning_engines

id = "projects/11092305173/locations/us-central1/reasoningEngines/8793735669278048256"
remote_app = reasoning_engines.ReasoningEngine(id)
remote_app.query(input="鬼切丸と氷の刃 フロストはどっちが強いですか?")

もし上手くいかない場合は、Cloud Logging をご確認ください。

6. 削除

不要となった Reasoning Engine の Resource を削除します。
先ほど作成した client.py に追記して、権限を持ったローカルの環境から実行してみます。

client.py
remote_app.delete()

まとめと補足

今回は Reasoning Engine から提供されている LangChainAgent クラスを利用して開発しましたが、すでに LangChain や他のフレームワークでアプリケーションを構築されている方は独自にカスタムテンプレートを利用して開発することも可能です。
さらに Python で記述できるということで、Streamlit や Google が開発している Mesop などの UI フレームワークと相性がいいのではないでしょうか。

なお今回は触れていませんが、 LangChain と Google Cloud のマネージドサービスとの連携も全面的に強化されており、近傍検索に利用するベクトルやチャット履歴の保存に Cloud SQL、Memorystore、Firestore などのマネージドデータベースを利用することができます。

LangChain on Vertex AI(プレビュー)は、現状では記載の通り各種制約がありますので、今後の改善と GA にご期待ください。

Google Cloud Japan

Discussion