pydantic-ai

from pydantic_ai import Agent

# 定义一个非常简单的智能代理，包括要使用的模型。你也可以在运行时设置模型。
agent = Agent(
    'gemini-1.5-flash',
    # 使用关键字参数注册一个静态系统提示。
    # 对于更复杂的动态生成提示，请参阅下例。
    system_prompt='Be concise, reply with one sentence.',
)

# 与大语言模型同步运行智能代理，进行对话。
# 这里的交流应该会很短：PydanticAI 会发送系统提示和用户查询给 LLM，
# 模型将返回一个文本响应。查看下方更复杂的运行示例。
result = agent.run_sync('Where does "hello world" come from?')
print(result.data)
"""
The first known use of "hello, world" was in a 1974 textbook about the C programming language.
"""

from dataclasses import dataclass

from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext

from bank_database import DatabaseConn


# SupportDependencies 用于传递数据、连接和逻辑到运行系统提示和工具函数时需要的模型中。
# 依赖注入为定制智能代理的行为提供了一种类型安全的方法。
@dataclass
class SupportDependencies:
    customer_id: int
    db: DatabaseConn


# 此 Pydantic 模型定义智能代理返回结果的结构。
class SupportResult(BaseModel):
    support_advice: str = Field(description='Advice returned to the customer')
    block_card: bool = Field(description="Whether to block the customer's card")
    risk: int = Field(description='Risk level of query', ge=0, le=10)


# 这个代理将成为银行的一线支持。
# 代理对于接受的依赖类型和返回的结果类型是通用的。
# 在这种情况下，支持代理的类型是 `Agent[SupportDependencies, SupportResult]`。
support_agent = Agent(
    'openai:gpt-4o',
    deps_type=SupportDependencies,
    # 代理的响应将被保证为 SupportResult，
    # 如果验证失败，代理将被提示重试。
    result_type=SupportResult,
    system_prompt=(
        'You are a support agent in our bank, give the '
        'customer support and judge the risk level of their query.'
    ),
)


# 动态系统提示可以利用依赖注入。
# 依赖通过 `RunContext` 参数传递，该参数由上面的 `deps_type` 参数化。
# 如果这里的类型注解错误，静态类型检查器将会捕获。
@support_agent.system_prompt
async def add_customer_name(ctx: RunContext[SupportDependencies]) -> str:
    customer_name = await ctx.deps.db.customer_name(id=ctx.deps.customer_id)
    return f"The customer's name is {customer_name!r}"


# `tool` 可以让你注册 LLM 在响应用户时可能调用的函数。
# 同样，依赖通过 `RunContext` 传递，其他任何参数都会成为工具传递给 LLM 的架构。
# Pydantic 用于验证这些参数，错误将被传回给 LLM，以便其重试。
@support_agent.tool
async def customer_balance(
    ctx: RunContext[SupportDependencies], include_pending: bool
) -> float:
    """Returns the customer's current account balance."""
    # 工具的文档字符串也会作为工具的描述传递给 LLM。
    # 参数描述从文档字符串中提取并加入到传递给 LLM 的参数架构中。
    balance = await ctx.deps.db.customer_balance(
        id=ctx.deps.customer_id,
        include_pending=include_pending,
    )
    return balance


...  # 在实际用例中，你会添加更多工具和更长的系统提示


async def main():
    deps = SupportDependencies(customer_id=123, db=DatabaseConn())
    # 异步运行智能代理，与LLM对话直到获得最终响应。
    # 即使在这种相对简单的情况下，代理也将与LLM交换多条消息，因为会调用工具来获取结果。
    result = await support_agent.run('What is my balance?', deps=deps)
    # 结果将由 Pydantic 验证以确保为 `SupportResult`，由于代理是通用的，
    # 它也将被类型化为 `SupportResult` 以帮助静态类型检查。
    print(result.data)
    """
    support_advice='Hello John, your current account balance, including pending transactions, is $123.45.' block_card=False risk=1
    """

    result = await support_agent.run('I just lost my card!', deps=deps)
    print(result.data)
    """
    support_advice="I'm sorry to hear that, John. We are temporarily blocking your card to prevent unauthorized transactions." block_card=True risk=8
    """