使用 Pydantic AI和千问为银行构建支持agent的一个小型但完整的示例

"""使用 Pydantic AI 为银行构建支持代理的一个小型但完整的示例。

使用以下命令运行：

    uv run main.py
"""

from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import List

from pydantic import BaseModel

from pydantic_ai import Agent, RunContext
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider

# 填入自己的key
qwen_key = 'sk-dxxxxxxxx'
base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
zhipu_model_str = 'qwen-flash'


class DatabaseConn:
    """这是一个用于示例目的的伪数据库。

    在实际应用中，您将连接到外部数据库
    （例如 PostgreSQL）以获取有关客户的信息。
    """
    _customers = {
        123: {'name': 'John', 'balance': 100.00, 'pending_balance': 123.45, 'address': '幸福路1号'},
        456: {'name': 'Jane', 'balance': 200.00, 'pending_balance': 250.00, 'address': '阳光大道2号'},
    }
    _transactions = {
        123: [
            {'date': (datetime.now() - timedelta(days=1)).isoformat(), 'description': '咖啡店消费', 'amount': -5.50},
            {'date': (datetime.now() - timedelta(days=2)).isoformat(), 'description': '超市购物', 'amount': -32.10},
            {'date': (datetime.now() - timedelta(days=3)).isoformat(), 'description': '工资入账', 'amount': 1500.00},
        ]
    }

    @classmethod
    async def customer_name(cls, *, id: int) -> str | None:
        """根据客户ID异步获取客户姓名。"""
        return cls._customers.get(id, {}).get('name')

    @classmethod
    async def customer_balance(cls, *, id: int, include_pending: bool) -> float:
        """根据客户ID异步获取客户余额。"""
        customer = cls._customers.get(id)
        if not customer:
            raise ValueError('未找到客户')
        return customer['pending_balance'] if include_pending else customer['balance']

    @classmethod
    async def get_recent_transactions(cls, *, id: int, limit: int = 5) -> List[dict]:
        """获取最近的交易记录。"""
        return cls._transactions.get(id, [])[:limit]

    @classmethod
    async def transfer_money(cls, *, from_id: int, to_id: int, amount: float) -> str:
        """在两个账户之间转账。"""
        if from_id not in cls._customers or to_id not in cls._customers:
            raise ValueError("无效的客户ID")
        if cls._customers[from_id]['balance'] < amount:
            raise ValueError("余额不足")
        
        cls._customers[from_id]['balance'] -= amount
        cls._customers[to_id]['balance'] += amount
        # 记录交易
        cls._transactions.setdefault(from_id, []).insert(0, {
            'date': datetime.now().isoformat(),
            'description': f'转账给客户 {to_id}',
            'amount': -amount,
        })
        return f"交易成功！从ID {from_id} 向ID {to_id} 转账 ${amount:.2f}。"

    @classmethod
    async def update_address(cls, *, id: int, new_address: str) -> bool:
        """更新客户的地址。"""
        if id in cls._customers:
            cls._customers[id]['address'] = new_address
            return True
        return False


@dataclass
class SupportDependencies:
    """定义支持代理运行时所需的依赖项。"""
    customer_id: int
    db: DatabaseConn = field(default_factory=DatabaseConn)


class SupportOutput(BaseModel):
    """定义支持代理的输出结构。"""
    support_advice: str
    """返回给客户的建议"""
    block_card: bool = False
    """是否冻结客户的卡"""
    risk: int
    """查询的风险等级"""



# 初始化大语言模型
model = OpenAIChatModel(
    zhipu_model_str,
    provider=OpenAIProvider(
        base_url=base_url, api_key=qwen_key
    ),
)

# 创建一个代理实例
support_agent = Agent(
    model,
    deps_type=SupportDependencies,
    output_type=SupportOutput,
    instructions=(
        '你是我行的支持代理，为客户提供支持并判断其查询的风险等级。'
        '请使用客户的名字进行回复。'
        '你可以查询余额、查询交易记录、执行转账、更新地址以及冻结卡片。'
        '对于一般性问题，如贷款服务，请提供通用信息。'
    ),
)


@support_agent.instructions
async def add_customer_name(ctx: RunContext[SupportDependencies]) -> str:
    """这是一个指令函数，用于在处理请求前获取并添加客户姓名到上下文中。"""
    customer_name = await ctx.deps.db.customer_name(id=ctx.deps.customer_id)
    return f"客户的名字是 {customer_name!r}"


@support_agent.tool
async def customer_balance(
    ctx: RunContext[SupportDependencies], include_pending: bool
) -> str:
    """这是一个工具函数，供AI代理调用以查询客户的当前账户余额。"""
    balance = await ctx.deps.db.customer_balance(
        id=ctx.deps.customer_id,
        include_pending=include_pending,
    )
    return f'${balance:.2f}'

@support_agent.tool
async def get_recent_transactions(
    ctx: RunContext[SupportDependencies], limit: int = 3
) -> List[dict]:
    """查询最近的交易记录。"""
    transactions = await ctx.deps.db.get_recent_transactions(
        id=ctx.deps.customer_id,
        limit=limit,
    )
    return transactions

@support_agent.tool
async def transfer_money(
    ctx: RunContext[SupportDependencies], to_customer_id: int, amount: float
) -> str:
    """执行转账操作。"""
    try:
        result = await ctx.deps.db.transfer_money(
            from_id=ctx.deps.customer_id,
            to_id=to_customer_id,
            amount=amount,
        )
        return result
    except ValueError as e:
        return str(e)

if __name__ == '__main__':
    # 创建依赖实例
    deps = SupportDependencies(customer_id=123)
    
    print("--- 场景1: 查询余额 ---")
    result = support_agent.run_sync('我的余额是多少？', deps=deps)
    print(result.output)
    # support_advice='您的当前余额是 $123.45，包括待处理的交易。' block_card=False risk=1
    print("\n" + "="*30 + "\n")

    print("--- 场景2: 丢失卡片 ---")
    result = support_agent.run_sync('我刚把卡丢了！', deps=deps)
    print(result.output)
    # support_advice='John，很抱歉听到您的卡丢了。为了保障您的账户安全，我建议立即冻结您的卡片。您可以放心，我会协助您完成此操作。' block_card=False risk=5
    print("\n" + "="*30 + "\n")

    print("--- 场景3: 查询最近交易 ---")
    result = support_agent.run_sync('我最近有什么交易记录吗？', deps=deps)
    print(result.output)
    # support_advice='以下是您最近的交易记录：\n1. 2025-11-03：咖啡店消费，金额 -5.50 元。\n2. 2025-11-02：超市购物，金额 -32.10 元。\n3. 2025-11-01：工资入账，金额 +1,500.00 元。' block_card=False risk=1
    print("\n" + "="*30 + "\n")

    print("--- 场景4: 执行转账 ---")
    result = support_agent.run_sync('帮我转50块钱给ID是456的用户。', deps=deps)
    print(result.output)
    # support_advice='已成功将 $50.00 转账给 ID 456。如有其他需求，请随时告知！' block_card=False risk=1
    print("\n" + "="*30 + "\n")

    print("--- 场景3.1: 查询最近交易 ---")
    result = support_agent.run_sync('我最近有什么交易记录吗？', deps=deps)
    print(result.output)
    # support_advice='以下是您最近的交易记录：\n- 2025-11-04：转账给客户 456，金额 -50.00 元\n- 2025-11-03：咖啡店消费，金额 -5.50 元\n- 2025-11-02：超市购物，金额 -32.10 元\n- 2025-11-01：工资入账，金额 +1500.00 元' block_card=False risk=1
    print("\n" + "="*30 + "\n")

    print("--- 场景5: 咨询银行服务 ---")
    result = support_agent.run_sync('你们银行有提供个人贷款服务吗？', deps=deps)
    print(result.output)
    # upport_advice='是的，我们银行提供个人贷款服务。您可以申请用于消费、装修、教育或旅行等用途的贷款。具体利率和还款期限根据您的信用状况和贷款金额而定。建议您联系我们的客户经理获取详细信息或在线提交申请。' block_card=False risk=1
    print("\n" + "="*30 + "\n")