🔖
PydanticAIでエージェントを作る-3:deps_typeの利用
すべてがdictにな〜れ〜
TL;DR
前回はPydantic AI[1]を使って、簡単なエージェント作成[2]、ツールの利用[3]をやってきました。今回は、入力の依存性(deps_type)の設定を使っていきたいと思います。
この機能、本当にいろいろなことが出来そうですが、今回は「スケジュールをプロンプトから更新するエージェント」を例にして、「入力に使ったデータを、エージェントの出力をもとに更新する」ケースを実装していきます。
- 入出力をPydanticのBaseModelで定義すると楽。
- BaseModel <-> dict の相互変換は楽。
(BaseModel.__dict__
<->BaseModelクラス(dict)
) - dictの部分更新は
dict.update(更新用dict)
でOK。- つまり、脳筋でdictを用意しておいて、要所要所でBaseModelに変換すれば入出力はなんとでもなる。
イメージ
コード全体
import difflib
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
from typing import List, Optional, Literal
from datetime import datetime
nest_asyncio.apply()
load_dotenv()
# 入力の型を定義
class schedule(BaseModel):
task: str # 入力の型を定義
Start: str = Field(
"",
description="The start time of the task. The format should be YYYY-MM-DD.",
) # 開始日時の型を定義
Finish: str = Field(
"",
description="The finish time of the task. The format should be YYYY-MM-DD.",
) # 終了日時の型を定義
class input_info(BaseModel):
prompts: str # 入力の型を定義
schedules: List[schedule]
# 入力の型を定義
class output_info(BaseModel):
schedules: List[schedule] # レスポンスの型を定義、入力の型と同じ
# レスポンスの型を定義
# Pydanticのエージェントを作成
# 環境変数にOPENAI_API_KEYは設定済み
scheduler_agent = Agent(
"openai:gpt-4o-mini", # モデルの指定
deps_type=input_info, # 入力の型を指定
result_type=output_info, # レスポンスの型を指定
system_prompt="You are an AI assistant helping a user with organizing schedules.",
)
@scheduler_agent.system_prompt
async def get_system_prompt(ctx: RunContext[input_info]) -> str:
"""
System prompt for the AI agent to handle with list_of_schedules.
"""
added_prompt = f"""
Toeday is {datetime.now().strftime("%Y-%m-%d")}.
User's current schedules:
{ctx.deps.schedules}.
"""
return added_prompt
if "user_input" not in st.session_state:
temp_input = input_info(prompts="", schedules=[]) # 入力の初期値を設定
st.session_state.user_input = (
temp_input.model_dump()
) # 入力をdictでセッションステートに保存
if "result" not in st.session_state:
temp_result = output_info(schedules=[]) # レスポンスの初期値を設定
st.session_state.result = (
temp_result.model_dump()
) # レスポンスをdictでセッションステートに保存
st.write("更新前データ")
st.json(st.session_state.user_input) # 入力をJSON形式で表示
st.table(
pd.DataFrame.from_dict(st.session_state.user_input["schedules"])
) # データフレームを表として表示
st.write("更新予定データ")
st.json(st.session_state.result) # レスポンスをJSON形式で表示
st.table(
pd.DataFrame.from_dict(st.session_state.result["schedules"])
) # データフレームを表として表示
# データ更新ボタン
if st.button("データ更新"):
st.session_state.user_input.update(st.session_state.result)
st.rerun()
# streamlitのUI作成
st.title("Pydanticを使ったAIエージェントの作成")
prompts = st.text_area("プロンプト入力欄", "ここに入力してください")
button1, button2 = st.columns(2)
if button1.button("チャット"):
st.write("チャットボタンがクリックされました")
user_input_deps = input_info(**st.session_state.user_input) # 入力を取得
user_input_deps.prompts = prompts # ユーザー入力を更新
result = scheduler_agent.run_sync(user_input_deps.prompts, deps=user_input_deps)
st.session_state.user_input.update(
user_input_deps.model_dump()
) # レスポンスをセッションステート
st.session_state.result.update(
result.data.model_dump()
) # レスポンスをセッションステートに保存
st.rerun()
とりあえず上記をstreamlitで動かしてみると、こんな感じになります。
プロンプトから指定した型どおりに出力ができています。
前回からの主な変化点
前回からの変化点としては、
- 入力用のclass
input_info
をPydanticのBaseModelで定義し、ここにpromptsとschedules(更新したい内容)を入れています。
class schedule(BaseModel):
task: str # 入力の型を定義
Start: str = Field(
"",
description="The start time of the task. The format should be YYYY-MM-DD.",
) # 開始日時の型を定義
Finish: str = Field(
"",
description="The finish time of the task. The format should be YYYY-MM-DD.",
) # 終了日時の型を定義
class input_info(BaseModel):
prompts: str # 入力の型を定義
schedules: List[schedule]
# 入力の型を定義
- 出力用のclass
output_info
も同様に定義しています。内部には、入力と同じ型のschedulesを持っています。
class output_info(BaseModel):
schedules: List[schedule] # レスポンスの型を定義、入力の型と同じ
- 入出力の型を指定する際に、
deps_type
とresult_type
を使っています。
scheduler_agent = Agent(
"openai:gpt-4o-mini", # モデルの指定
deps_type=input_info, # 入力の型を指定
result_type=output_info, # レスポンスの型を指定
system_prompt="You are an AI assistant helping a user with organizing schedules.",
)
- 入力のschedulesを個別にsystem_promptに表示するために、
get_system_prompt
を定義しています。これで、エージェントのシステムプロンプトには、入力のschedulesと今日の日付が追加されます。 - 前回軽く触れた
RunContext
を使って、入力のschedulesを取得しています。
ctxの中にinput_infoから取得した情報が入っており、ctx.deps.~~
で取得できます。(今回はctx.deps.schedules
)
@scheduler_agent.system_prompt
async def get_system_prompt(ctx: RunContext[input_info]) -> str:
"""
System prompt for the AI agent to handle with list_of_schedules.
"""
added_prompt = f"""
Toeday is {datetime.now().strftime("%Y-%m-%d")}.
User's current schedules:
{ctx.deps.schedules}.
"""
return added_prompt
- 出力のschedulesを、入力のschedulesに更新できるボタンを作っています。
なお、各データの一時保存にはst.session_stateを使っています。
if st.button("データ更新"):
st.session_state.user_input.update(st.session_state.result)
st.rerun()
- データの更新は
- PydanticのBaseModelとなっている出力を
model_dump()
でdictに変換し、 - 各dictを
update()
で更新しています。
- PydanticのBaseModelとなっている出力を
user_input_deps = input_info(**st.session_state.user_input) # 入力を取得
user_input_deps.prompts = prompts # ユーザー入力を更新
result = scheduler_agent.run_sync(user_input_deps.prompts, deps=user_input_deps)
st.session_state.user_input.update(
user_input_deps.model_dump()
) # レスポンスをセッションステート
st.session_state.result.update(
result.data.model_dump()
) # レスポンスをセッションステートに保存
ポイント
- 入出力の共通部分(更新部分)を同一の型で定義すると、更新が楽。
- 今回はschedulesだけでしたが、更新部分の数が増えても入出力の型管理だけ正しく行えば、
dict~~.update()
ですべてOK。(アイテムごとに記述する必要なし!) - 気づけばこの記事、結局「
dict
美味しいです」しか言ってない。。。
Discussion