🔖

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

2025/01/24に公開

すべてが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_typeresult_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()で更新しています。
    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美味しいです」しか言ってない。。。

リンク

脚注
  1. Pydantic AI ↩︎

  2. PydanticAIが有能すぎて今までが何だったんだ感 ↩︎

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

Discussion