🔖

PydanticAIでエージェントを作る-2:toolの利用

2025/01/19に公開

TL;DR

前回はPydantic AI[1]を使って、簡単なエージェントを作成しました。
今回は、エージェントにtoolを追加して、簡単な計算を行います。

  • @"agent名".toolデコレータを使って、toolを定義します。
  • tool内のdocstringに説明を書けば、エージェントのシステムプロンプトにはtoolの説明を書かなくても良い模様。
  • toolの引数もエージェントが勝手に見てくれているようで、こちらもシステムプロンプトには記載不要。
    (docstringに型指定を書かなくてもに正しく対応できています。)
  • つまり、toolを後付しても、docstringに仕様さえ書けば、もともとのエージェントの設定をいじくらなくていい。。。
    マジか。。。

イメージ

コード全体

import nest_asyncio
import pandas as pd
import streamlit as st
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext

nest_asyncio.apply()
load_dotenv()

# Pydanticのエージェントを作成
# 環境変数にOPENAI_API_KEYは設定済み
pd_agent = Agent(
    "openai:gpt-4o-mini",  # モデルの指定
    system_prompt="You are an AI assistant helping a user with general questions.",
)


# ツールの定義

# べき乗の計算用のツールを定義
# ctx: RunContextはDependencyの受皿などに使うが、今回は使わないので型だけ指定(ないとエラー扱いになる。)
@pd_agent.tool
async def calculate_power(ctx: RunContext[str], a: float, b: float) -> float:
    """
    Tool to get A to the power of B.
    """
    c = a**b
    return c


# レスポンスの型を定義
class Structured_Out(BaseModel):  # レスポンスの型を定義
    response: str
    answer: float = Field(..., description="The result of the calculation.")
    used_tools: str = Field(..., description="The tools used in the calculation.")


# streamlitのUI作成
st.title("Pydanticを使ったAIエージェントの作成")

prompts = st.text_area("プロンプト入力欄", "ここに入力してください")

button1, button2 = st.columns(2)
if button1.button("チャット"):
    st.write("チャットボタンがクリックされました")
    result = pd_agent.run_sync(prompts, result_type=Structured_Out)
    result_dict = result.data.model_dump()  # レスポンスを辞書型に変換
    st.json(result_dict)  # レスポンスをJSON形式で表示
    df = pd.DataFrame([result_dict])  # データフレームに変換
    # データフレームを表として表示
    st.table(df)
    st.json(result.all_messages())

とりあえず上記をstreamlitで動かしてみると、こんな感じになります。

プロンプトでは「2.3の3.5乗」という生成AI単身では厳しそうな計算をお願いしていますが、エージェントがtoolに渡してちゃんと計算してくれています。

前回からの主な変化点

前回からの変化点としては、

  • べき乗計算用にtoolを定義しました。と言っても@pd_agent.toolデコレータをつけて関数定義だけです。
    なお、最初の引数としてRunContextを1回入れないと怒られるので、ここでは意味のない引数を入れています。
@pd_agent.tool
async def calculate_power(ctx: RunContext[str], a: float, b: float) -> float:
    """
    Tool to get A to the power of B.
    """
    c = a**b
    return c
  • 回答の型を以下のように定義しました。
    実際の回答を見ると正しく反映されていることがわかります。
class Structured_Out(BaseModel):  # レスポンスの型を定義
    response: str
    answer: float = Field(..., description="The result of the calculation.")
    used_tools: str = Field(..., description="The tools used in the calculation.")
  • toolへの受け渡しを見るために、エージェントのリクエスト・レスポンスの履歴を表示するようにしました。
    ちゃんと、toolに正しい引数の型で渡していることがわかります。
    st.json(result.all_messages())

ポイント

  • toolの追加は、@~~~~.toolデコレータを使って、関数を定義するだけ。
  • toolの最初の引数は(おまじないとして)RunContextを入れる。
  • その他の引数は、型を指定しておけば、エージェントが勝手に見てくれる。
  • toolのdocstringに説明を書けば、エージェントのシステムプロンプトにはtoolの説明を書かなくても良い。

RunContextについては、今回はtoolの性質上まともに使いませんでしたが、エージェントにユーザープロンプト以外の複数の変数(例えば、文書、RAGへの接続情報など)を渡したいときに非常に重宝します。
今後の記事で、そのあたりを書きたいと思います。

リンク

脚注
  1. Pydantic AI ↩︎

Discussion