Pydantic AI

  • Pydantic AI とはLLM Agentを構築するためのライブラリ
  • Pydantic modelを使ってTool Callのスキーマの定義ができる

Tool Call

from pydantic_ai import Agent
 
agent = Agent(
    'openai:gpt-4o',
    system_prompt="計算をして下さい",
)
 
@agent.tool_plain
async def add(a: int, b: int) -> int:
    """`a` と `b` を足し算します。
 
    Args:
        a: 1つ目の数値
        b: 2つ目の数値
    """
    return a + b
 
async def main():
    result = await agent.run('1 + 2?')
    print(result.data)

Result Type

LLMの返す値やToolの引数・返り値型をPydantic modelで指定できる。

from pydantic import BaseModel, Field
 
class Point(BaseModel):
    x: int = Field(description="x座標")
    y: int = Field(description="y座標")
 
agent = Agent(
    'openai:gpt-4o',
    system_prompt="計算をして下さい",
    result_type=Point,  # ←
)
 
@agent.tool_plain
async def add(a: Point, b: Point) -> Point:  # ←
    ...
 
async def main():
    result = await agent.run('(1, 2) + (3, 4)?')
    answer: Point = result.data  # ←

Deps

  • Tool Call時に依存関係を指定できる
    • 固定値やデータベースのコネクションを持たせることが出来る
  • 依存を使って動的なPromptを生成することも可能
from pydantic_ai import Agent, RunContext
from dataclasses import dataclass
 
@dataclass
class Deps:
    factor: int
 
agent = Agent(
    'openai:gpt-4o',
    deps_type=Deps,  # ←
)
 
@agent.system_prompt
async def dynamic_prompt(ctx: RunContext[Deps]) -> str:
    return f"計算をして最後に {ctx.deps.factor} を掛けて下さい"
 
@agent.tool
async def add(ctx: RunContext[Deps], a: int, b: int) -> int:
    ...
 
async def main():
    deps = Deps(factor=3)
    result = await agent.run('1 + 2?', deps=deps)  # ←
    print(result.data)

Agentのネスト

agent_a = Agent(
    'openai:gpt-4o',
    system_prompt="...",
)
 
agent_b = Agent(
    'openai:gpt-4o',
    system_prompt="...",
)
 
@agent_a.tool
async def tool() -> str:
    result = await agent_b.run("...")
    ...
 
async def main():
    result = await agent_a.run("...")
    ...

動的なToolの使用

Pydantic ModelはModel → JSON Schema etc. の変換は出来るが逆は出来ない。
datamodel-code-generator はJSON SchemaからModelのコード(文字列)を返すことしか出来ず eval()exec() が必要になってしまう。
また、Pydantic自体にも create_model({ "hoge": (str, Field(...), ...) }) のように動的にモデルを作る方法はあるが、任意のJSON Schemaを走査して create_model に渡すdictを作る処理を書く必要がある。

幸い、Pydantic AIにPydantic Modelを介さず直接JSON Schemaを渡してToolを作成する方法があったので、それを使う。

def post_mail(ctx: RunContext[Deps], body: dict[str, Any]) -> str:  # ← body は dict のまま受ける
    """メールを投稿します。
 
    Args:
        body: メールの内容
    """
    print(body)
    return "ok"
 
async def prepare_post_mail(
    ctx: RunContext[Deps], tool_def: ToolDefinition
) -> ToolDefinition | None:
    corner = ctx.deps.client.corner(ctx.deps.corner_id)
    tool_def.description = corner.description
    tool_def.parameters_json_schema = {  # ← JSON Schema を直接差し込む
        "type": "object",
        "required": ["body"],
        "properties": {
            "body": corner.mail_schema,
        }
    }
    return tool_def
 
post_mail = Tool(post_mail, prepare=prepare_post_mail)  # ← prepare で実行時にschemaを決める
 
agent = Agent(
    "openai:gpt-4o-mini",
    system_prompt="""
    あなたはコーナーにメールを投稿するリスナーです。
    """,
    tools=[post_mail],
)

環境構築メモ

sudo apt-get install libpq-dev
poetry add sqlalchemy[asyncio] psycopg2

参考文献