150 lines
4.2 KiB
Python
150 lines
4.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
银行客户服务
|
|
"""
|
|
from typing import Any, AsyncGenerator
|
|
|
|
from pydantic import BaseModel, Field
|
|
from pydantic_ai import (
|
|
Agent,
|
|
Capability,
|
|
DeferredToolResults,
|
|
ModelMessage,
|
|
RunContext,
|
|
RunUsage,
|
|
UsageLimits,
|
|
)
|
|
from pydantic_ai.messages import (
|
|
AgentStreamEvent,
|
|
PartStartEvent,
|
|
ToolCallPart,
|
|
FunctionToolResultEvent,
|
|
ToolReturnPart,
|
|
)
|
|
from pydantic_ai.run import AgentRunResultEvent, AgentRunResult
|
|
|
|
from application.workshop.models import (
|
|
DEEPSEEK_V4_FLASH_MODEL,
|
|
DEEPSEEK_V4_FLASH_MODEL_SETTINGS,
|
|
)
|
|
from asyncpg import Pool
|
|
from application.workshop.events import AgentRuntimeEvent
|
|
|
|
|
|
class Deps(BaseModel):
|
|
"""
|
|
依赖类
|
|
"""
|
|
|
|
model_config = {"arbitrary_types_allowed": True} # 允许任意类型
|
|
|
|
account_id: str = Field(..., description="账户唯一标识")
|
|
connection_pool: Pool = Field(..., description="连接池")
|
|
|
|
|
|
class Output(BaseModel):
|
|
"""
|
|
输出类
|
|
"""
|
|
|
|
response: str = Field(..., description="客服应答")
|
|
|
|
|
|
# 账户上下文能力
|
|
account_context = Capability[Deps](
|
|
id="account-context",
|
|
description="提供账户信息，用于客服应答",
|
|
defer_loading=False, # 延迟加载置为否，立即加载
|
|
)
|
|
|
|
|
|
@account_context.instructions
|
|
async def query_account_name(ctx: RunContext[Deps]) -> str:
|
|
"""
|
|
查询账户户名
|
|
"""
|
|
async with ctx.deps.connection_pool.acquire() as conn:
|
|
account_name = await conn.fetchval(
|
|
"SELECT account_name FROM bank_customer_service WHERE account_id = $1",
|
|
ctx.deps.account_id,
|
|
)
|
|
if account_name is None:
|
|
return "该账户不存在"
|
|
return f"该账户户名是 {account_name!r}"
|
|
|
|
|
|
@account_context.tool(name="查询账户余额")
|
|
async def query_account_balance(ctx: RunContext[Deps]) -> str:
|
|
"""
|
|
查询账户余额
|
|
"""
|
|
async with ctx.deps.connection_pool.acquire() as conn:
|
|
account_balance = await conn.fetchval(
|
|
"SELECT account_balance FROM bank_customer_service WHERE account_id = $1",
|
|
ctx.deps.account_id,
|
|
)
|
|
if account_balance is None:
|
|
return "该账户不存在"
|
|
return f"该账户余额： {account_balance:.2f}"
|
|
|
|
|
|
# 智能体
|
|
agent = Agent[Deps, Output](
|
|
model=DEEPSEEK_V4_FLASH_MODEL,
|
|
model_settings=DEEPSEEK_V4_FLASH_MODEL_SETTINGS,
|
|
deps_type=Deps,
|
|
output_type=Output,
|
|
instructions="你是我方银行的客服智能体，为客户提供客服应答。",
|
|
capabilities=[account_context],
|
|
)
|
|
|
|
|
|
async def run_stream_events(
|
|
deps: Any,
|
|
user_prompt: str | None,
|
|
message_history: list[ModelMessage],
|
|
usage: RunUsage,
|
|
usage_limits: UsageLimits | None = None,
|
|
deferred_tool_results: DeferredToolResults | None = None,
|
|
) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent[Output] | None,]:
|
|
"""
|
|
运行并流式输出事件
|
|
"""
|
|
if not isinstance(deps, Deps):
|
|
deps = Deps.model_validate(deps)
|
|
|
|
async with agent.run_stream_events(
|
|
user_prompt=user_prompt,
|
|
deps=deps,
|
|
message_history=message_history,
|
|
usage=usage,
|
|
usage_limits=usage_limits,
|
|
deferred_tool_results=deferred_tool_results,
|
|
) as events:
|
|
async for event in events:
|
|
match event:
|
|
case PartStartEvent(
|
|
part=part,
|
|
):
|
|
match part:
|
|
case ToolCallPart(tool_name=tool_name):
|
|
if tool_name.startswith("final_result"):
|
|
continue
|
|
|
|
case FunctionToolResultEvent(part):
|
|
match part:
|
|
case ToolReturnPart(
|
|
tool_name=tool_name,
|
|
):
|
|
if tool_name.startswith("final_result"):
|
|
continue
|
|
|
|
case AgentRunResultEvent(result=result):
|
|
match result:
|
|
case AgentRunResult(output=output):
|
|
match output:
|
|
case Output(response=response):
|
|
yield AgentRuntimeEvent(content=response)
|
|
|
|
yield event
|