🔖
PydanticAIでエージェントを作る-2:toolの利用
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への接続情報など)を渡したいときに非常に重宝します。
今後の記事で、そのあたりを書きたいと思います。
Discussion