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の例
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