【AI Shift Advent Calendar 2024】PydanticAI で実現する Dependency Injection

はじめに

こんにちは、AI Shift のAIチームに在籍している長澤 (@sp_1999N) です!

この記事は AI Shift Advent Calendar 2024 の20日目の記事になります。

今回は比較的新しく登場した PydanticAI を使って、LLM の実運用におけるDependency Injection (DI, 依存性の注入) を検討してみたいと思います。

DI とは何かというトピックについても簡単に解説していますので、LLM を使用したアプリケーションを開発している方はどなたでもお読みいただけるものになっているかと思います!

(この記事の内容は2024/12時点のものになります。PydanticAI は現時点で開発が盛んなようなので、最新の情報は公式ドキュメントをご参照いただければと思います。)

PydanticAI とは

PydanticAI は Python で利用できる Agent Framework です。Python で型安全な実装を行おうとした際に利用される Pydantic の系譜を受け継いだものになっています。

Similarly, virtually every agent framework and LLM library in Python uses Pydantic, yet when we began to use LLMs in Pydantic Logfire, we couldn't find anything that gave us the same feeling.

We built PydanticAI with one simple aim: to bring that FastAPI feeling to GenAI app development.

こちらのコメントにもあるように、「FastAPI での直感的で快適な使用感を生成AIアプリケーションにもたらすこと」を目的として開発されています。API で利用可能な LLM をプロダクションレベルで運用する際の利用が想定されています。

PydanticAI の特徴としては以下のものが挙げられています。

  • Built by the team behind Pydantic (the validation layer of the OpenAI SDK, the Anthropic SDK, LangChain, LlamaIndex, AutoGPT, Transformers, CrewAI, Instructor and many more)
  • Model-agnostic — currently OpenAI, Anthropic, Gemini, Ollama, Groq, and Mistral are supported, and there is a simple interface to implement support for other models.
  • Type-safe
  • Control flow and agent composition is done with vanilla Python, allowing you to make use of the same Python development best practices you'd use in any other (non-AI) project
  • Structured response validation with Pydantic
  • Streamed responses, including validation of streamed structured responses with Pydantic
  • Novel, type-safe dependency injection system, useful for testing and eval-driven iterative development
  • Logfire integration for debugging and monitoring the performance and general behavior of your LLM-powered application

紹介されている特徴を見ると、Pydantic の恩恵を受けた型安全な開発を LLM provider に依存せずに実装できることが伺えます。

また従来の Pydantic と比較して、dependency injection system が新しい機能として紹介されています。

今回の記事ではこの DI 機能を紐解きながら、LLM の実運用における DI の実装例をご紹介します。

Dependency Injection とは

本題に突入する前に、DI についても簡単に触れておきます。

ソフトウェア開発の経験が豊富な場合は聞き馴染みのある言葉かも知れませんが、そうでない場合はなかなかにとっつきにくいトピックだと思います。(既に馴染みある方は次のセクションに進んでいただいて問題ありません。)

Dependency Injection とは文字通り「プログラミングにおける依存性を (外部から) 注入する分離可能な形での管理」を実現するためのものです。うまく依存性を切り離せると、単体機能の管理・テストが容易になるなどの嬉しさがあります。ただこれだけでは分かりにくいので、具体例を交えながら解説します。

例えば、LLM でのメール校正・自動返信サービスを作るような場合を考えます。

全ての処理を分離せずに書くことによる課題

この時、全ての処理を1つの関数で実装したプログラムを作成しても良いかも知れませんが、どの機能がきちんと動いているのかのテストが大変になります。

また「校正処理を LLM ではなくルールベースのロジックに置き換えたい」や「送信処理で使用していたライブラリなどを大きく変更したい」などの場面を考えても、上記のような場合では改修作業は骨が折れそうな予感がします。

モジュールへの分割

そこで次に出てくるのがモジュールベースでの実装になります。ここで考えるサービスは例えば「メールの受信・前処理」「LLM によるメールの校正」「校正内容を反映したメールの自動返信」などに分けることができます。

分離したモジュールごとにそれぞれの責任範囲を定義することで、全体の構造がより理解しやすく、保守しやすい設計になります。このような設計は 単一責任原則 (Single Responsibility Principle) に基づいています。つまり、各モジュールやコンポーネントはその役割を1つのことに限定し、それぞれが具体的なタスクを担当します。

モジュール分割した後に残る課題

しかし、モジュールベースで実装した場合でも、各モジュールの「依存関係」の管理が課題として残ります。たとえば、「LLM によるメール校正」モジュールが特定のAPIクライアントに直接依存していると、将来的に校正ロジックを別のコンポーネントに置き換えたい場合や、異なるAPIに切り替えたい場合に、依存コードを大幅に書き直す必要が出てきます。

このような課題に対して、Dependency Injection が役に立ちます。

DI の簡単な例

例えば上記の例について DI の思想に基づいた実装例とそうでない例を簡単に比較してみます。

まずは全ての処理をまとめた例になります。上述していたモジュールが全て1つの関数に集約されてしまっています。

class EmailProcessingService:
    def process_email(self, email):
        prompt = "Please proofread the given email."
        # 特定の LLM ライブラリに依存
        llm = SpecificLLMLibrary()
        corrected_email = llm.generate_response(prompt, email)

        # 特定のSMTPライブラリに依存
        smtp_client = SpecificSMTPClient()
        smtp_client.send(corrected_email)

メタ的に process_email 関数を見ると、この関数は「メールを入力したら自動で校正されて送信されるサービス」としての責務を果たしてくれれば十分であり、どのロジックでメールが校正されているかは関心の外側 (また別の話) になります。

しかしこのままだと、例えば校正処理をルールベースのロジックに変更しようとすると process_email の関数を大きく改修する必要があります。

ここで、依存性を簡単に注入した例を見てみます。

class EmailProcessingService:
    def __init__(self, llm, prompt, mail_client):
        # 依存性を外部から注入
        self.llm = llm
        self.prompt = prompt
        self.mail_client = mail_client

    def process_email(self, email):
        corrected_email = self.llm.generate_response(self.prompt, email)
        self.mail_client.send(corrected_email)

各モジュールを引数レベルに引き上げただけに見えるかも知れませんが、process_email の関数内で LLM が何のモデルであるか、メール送信に使っているライブラリが何であるかを意識せずに良くなっていることが分かります。(ただし、それぞれのモジュールにおける特定の処理の呼び出し方法 = インターフェースは他で定義しておく必要はあります。)

そしてさらにこの関数のテストケースを考えてみます。

# Mock の定義
class MockLLM:
    def generate_response(self, prompt, email):
        return f"Prompt: {prompt}, Corrected: {email}"

class MockMailClient:
    def send(self, email):
        print(f"Mock send: {email}")

# テスト用の依存性注入
mock_llm = MockLLM()
mock_mail_client = MockMailClient()
email_service = EmailProcessingService(mock_llm, mock_mail_client)

# テストの実行
email_service.process_email("Test email")

LLM やメール送信部分を簡単にモックとして定義しているだけですが process_email 関数が動くかどうかのテストができるようになっています。一番最初の例ではそれぞれのモジュール全てが完成しないとテストできない (依存度が高い) 状態でしたが、うまく依存性を切り離すことで単体機能の管理やテスト等がしやすくなっていることが分かります。

前置きがかなり長くなってしまいましたが、以上が Dependency Injection の簡単なご紹介になります。

PydanticAI の基本概念

PydanticAI における DI の説明のため、まずは PydanticAI における基本を押さえておきます。

PydanticAI では Agent が LLM とのインタラクションにおける主要なインターフェースになります。

そしてこの Agent を中心に以下のコンポーネントを使用します。(言い換えると Agent は以下の要素のコンテナとして機能します。)

  • system prompt
    • LLM に対するいわゆるシステムプロンプト
    • インスタンス化の際や、デコレータとして後からの引き渡しが可能です
  • function tool
    • LLM が外部情報等にアクセスするために呼び出すツール (関数)
    • インスタンス化の際や、デコレータとして後からの引き渡しが可能です
  • result type
    • LLM のレスポンス型を定義したもので、インスタンス化の時のみ指定可能
  • dependencies
    • 上記の system_prompt, tool そして result validators (result type に対する validation) のそれぞれについて Agent からのアクセスを提供します

上記のシステムプロンプトやツールは DI を利用して動的に設定することも可能ですが、インスタンス化の際に指定することも可能です。Agent のコンストラクタは以下のとおりです。

__init__(
    model: Model | KnownModelName | None = None,
    *,
    result_type: type[ResultData] = str,
    system_prompt: str | Sequence[str] = (),
    deps_type: type[AgentDeps] = NoneType,
    name: str | None = None,
    model_settings: ModelSettings | None = None,
    retries: int = 1,
    result_tool_name: str = "final_result",
    result_tool_description: str | None = None,
    result_retries: int | None = None,
    tools: Sequence[
        Tool[AgentDeps] | ToolFuncEither[AgentDeps, ...]
    ] = (),
    defer_model_check: bool = False,
    end_strategy: EndStrategy = "early"
)

PydanticAI での DI

それでは本題に入っていきます。繰り返しになりますが、PydanticAI では dependency injection system として system prompts, tools そして result validators の大きく3つが提供されています。

つまり、システムプロンプトやツール、返り値の型設定などを依存関係があるものとして外部から注入できるようになっています。

PydanticAI において依存関係にあるものは、RunContext 型を通して Agent からアクセスされます。

この解説のため、ドキュメントに掲載されている例 (dice_game.py) を少し改変して説明します。(改変した部分はコメントとして明記します。)

import os
import random
import argparse
from dotenv import load_dotenv
from pydantic_ai import Agent, RunContext

load_dotenv()   # 改変: 環境変数の設定


# Define the AI agent with the game logic
agent = Agent(
    'openai:gpt-4o',
    deps_type=str,   # DIの型を指定(今回はstr)
    system_prompt=(
        "You're a dice game, you should roll the die and see if the number "
        "you get back matches the user's guess. If so, tell them they're a winner. "
        "Use the player's name in the response."
    ),
)

@agent.tool_plain  
def roll_dice() -> str:
    """Roll a six-sided die and return the result."""
    return str(random.randint(1, 6))

@agent.tool  
def get_player_name(ctx: RunContext[str]) -> str:
    """Get the player's name."""
    return ctx.deps

def main():
    parser = argparse.ArgumentParser(description="Dice Guessing Game")
    parser.add_argument("--guess", type=int, help="Your guess (number between 1 and 6)")
    parser.add_argument("--name", type=str, required=True, help="Your name")
    args = parser.parse_args()

    # Agentの実行: DIとしてargsで受け取った値を注入
    dice_result = agent.run_sync(f'My guess is {args.guess}', deps=args.name)
    print(dice_result.data)

if __name__ == "__main__":
    main()
> python dice_game.py --guess 3 --name "Ben"
The roll came up 5. Sorry, Ben, your guess was 3. Better luck next time!

この例は、ダイスロールの出目を予想するプログラムになります。ここで注目して頂きたいのは2つの関数 roll_dice()get_player_name(ctx: RunContext[str])です。

Agent を定義した後で2つのツールを提供しているのですが、前者は @agent.tool_plain (依存性の注入が不要なツールの引き渡し)、後者は @agent.tool (依存性の注入が必要なツールの引き渡し) のデコレータを使用しており、かつ後者は引数に ctx: RunContext[str]を設定しています。

roll_dice()の関数はランダムな数値を返却するのみになっており、特定の外部情報は必要としていません。これと比べ get_player_name(ctx: RunContext[str]) はプログラム引数で渡される name の情報が必要になります。このため、引数に RunContext を設定し依存性を注入しています。

混乱を避けるため、ここまでの内容を整理します。

@agent.hoge のデコレータを使用することで Agent が利用する hoge (tool や prompt) を設定することができます。

もしこの設定において動的な処理 (依存性の注入が必要な場合) は RunContext を使って処理する形になります。逆に静的な処理 (依存性の注入が不要な場合) であれば RunContext を使用する必要はありません。

Agent 自体を DI する

もう1つ面白い例として、Agent 自体の DI をご紹介します。

PydanticAI の dependencies では python の任意の型を指定できます。つまり、Agent 型も DI として設定できる値になります。

実際に Agents as dependencies of other Agents のセクションでこの具体例が紹介されています。

ここではこのプログラムを参考に、独自の "冗談-1 グランプリ" を開催してみようと思います。

審査員として定義する judge_agent が主導してプログラムを実行します。

具体的には、ユーザーから冗談-1グランプリの開催の依頼を受け、審査員 Agent (judge_agent)がお笑い芸人 Agent (factory_agent)に冗談を生成させます。

そして芸人 Agent が生成した冗談に対して、審査員 Agent に評価をしてもらおうと思います。

from dotenv import load_dotenv
from dataclasses import dataclass
from pydantic_ai import Agent, RunContext


load_dotenv()

# AgentをDIするための型を定義
@dataclass
class MultiAgentDeps:
    factory_agent: Agent[None, list[str]]


judge_agent = Agent(
    'openai:gpt-4o',
    deps_type=MultiAgentDeps,
    system_prompt=(
        '"joke_factory" を使用して複数の冗談を生成してください。'
        'その後で冗談-1グランプリの審査員として、それぞれの冗談に対して順位とコメントを付けてください。'
    ),
)

# 仮想のお笑い芸人として振る舞うAgentを作成
factory_agent = Agent('gemini-1.5-pro',
                      result_type=list[str],
                      )

# お笑い芸人に冗談を生成させるための関数
@judge_agent.tool
async def joke_factory(ctx: RunContext[MultiAgentDeps], count: int) -> str:
    print(f"{count=}")
    r = await ctx.deps.factory_agent.run(f'冗談を{count}個生成してください。')
    return '\n'.join(r.data)

# 審査員は、グランプリの開催依頼をもとに、芸人に冗談を生成させる
result = judge_agent.run_sync('冗談-1グランプリの開催をお願いします。',
                             deps=MultiAgentDeps(factory_agent))
print(result.data)
> python depends_agent.py               
count=3
冗談を以下に示しますので、それぞれに順位とコメントを付けます。

1位: 「なぜ海賊はカレンダーを信用できないのですか? 彼らはいつも日付を決めようとしているからです。」
   - コメント: 海賊と「日付を決める」ことの二重の意味を使ったワードプレイがうまいですね。シンプルながらもクスッと笑える冗談です。

2位: 「スプーンとフォークのどちらが寂しいですか? フォークです。スプーンは全部のスープにありえるからです。」
   - コメント: スプーンとフォークという日常アイテムを使い、さりげない観察からくるユーモアが魅力的です。スプーンの人気を引き合いに出したアイデアが面白いです。

3位: 「ドアノブを壊して投獄されました。それは変えられないと思いました。」
   - コメント: ドアノブを壊すことで「変えられない」という状況に対するアイロニーが少し分かりにくく、他の冗談に比べてややインパクトに欠けた印象です。

以上です。楽しんでいただけましたでしょうか?

実行結果を見ると、factory_agent の生成結果が DI され、それに対する審査結果が無事に生成されていることが分かります。(まだまだお笑いは難しいようですね。)

ちなみに count の値を明に指定してはいませんが、今回は 3 として実行されています。この挙動が自明ではないので念のため説明します。

プログラムでは judge_agent が主体となっています。ツールとして joke_factory 関数をいわゆる function calling 的に呼び出します。そしてこの関数の中で、factory_agent が呼び出され、冗談を生成結果として関数の返り値に設定します。この時、joke_factory 関数では、冗談をいくつ生成するかの指定も可能です。この count の指定は joke_factory の関数を呼び出している judge_agent が指定しています。

試しに judge_agent のプロンプトに「冗談は4つでお願いします」と指定の文言を入れると以下のような挙動が観察されました。

> python depends_agent.py
count=1
count=1
count=1
count=1
冗談-1グランプリのランキングとコメントは以下の通りです。

### 第1位
**「なぜ海賊はそんなに悪い成績だったの? 彼らはいつもCを得ていたからです。」**

コメント: 「C」と「海(Sea)」をかけたユーモアが秀逸です。簡潔で理解しやすく、爆笑を生む最高の冗談です。

### 第2位
**「なぜ海賊は自分の船を停泊できないのですか? 彼らは停泊する方法を知らないからです!」**

コメント: 停泊(停める)という動作にちなんだシンプルなジョークがツボにはまります。若干わかりづらい印象もありますが、考えると面白いです。

### 第3位
**「なぜ海賊はそんなに悪い成績だったの? 彼らはいつも略奪していたから!」**

コメント: 海賊の「略奪」を成績に絡めたアイデアが斬新ですが、1位の冗談と似ているため、新鮮さに欠けます。

### 第4位
**「なぜ海賊はそんなに下手なカードプレイヤーなのですか?」**

コメント: やや曖昧で特にオチが感じられないため、順位が上がりませんでした。カードプレイの部分にユーモアを盛り込むとさらに面白いかもしれません。

挙動を観察する限り、今回のケースでは指定した冗談の数になるように複数回 joke_factory 関数を呼び出していることが分かります。(count=1 が複数回標準出力されていることから、1回ずつ冗談を生成させていることが分かります。)

DI により、function calling がシームレスに実行されているとも考えられるため、よりシンプルな実装が実現できています。

まとめると、上記のように DI をうまく使用することで、複数の LLM を組み合わせた、いわゆるマルチエージェントについてもメンテナンスしやすい形で実装することができました。

おわりに

Pydantic ひいては PydanticAI は型安全な実装を提供してくれることが分かります。

一方で RAG などをはじめとした応用事例からも分かるように、どのような情報を LLM に提供するか / アクセス可能な手段を LLM に提供するかは重要なトピックです。

しかしこれらの情報アクセスの設定はユースケースによって様々です。つまり適切に依存関係を切り分けることが保守・運用において重要な取り組みとなります。

このような現場における課題に対し、PydanticAI の dependency injection system は有用なアプローチになる気がしました。

今回は触れませんでしたが、PydanticAI ではチャット履歴の管理や Agent のテストについても effort less な実装を提供してくれています。

またストリーミング出力に構造を持たせた streaming structured response などの魅力的な提供もあります。

LLM をプロダクトに最適な形で組み込めるよう、AI チームとしても引き続きこの辺りの情報にキャッチアップしていければと考えております。

ここまでお読み頂きありがとうございました。

PICK UP

TAG