157 lines
7.5 KiB
Python
157 lines
7.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
银行客户服务
|
||
"""
|
||
from typing import AsyncGenerator
|
||
|
||
from pydantic import BaseModel, Field
|
||
from pydantic_ai import (
|
||
Agent,
|
||
Capability,
|
||
DeferredToolResults,
|
||
ModelMessage,
|
||
RunContext,
|
||
RunUsage,
|
||
UsageLimits,
|
||
)
|
||
from pydantic_ai.messages import AgentStreamEvent
|
||
from pydantic_ai.run import AgentRunResultEvent
|
||
|
||
from application.workshop.database import connection_pool
|
||
from application.workshop.models import (
|
||
DEEPSEEK_V4_FLASH_CHAT_MODEL,
|
||
DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS,
|
||
)
|
||
|
||
|
||
class Deps(BaseModel):
|
||
"""
|
||
依赖类
|
||
"""
|
||
|
||
user_id: str = Field(..., description="用户唯一标识")
|
||
|
||
|
||
# 用户上下文能力
|
||
user_context = Capability[Deps](
|
||
id="user-context",
|
||
description="提供用户信息，用于客服应答",
|
||
defer_loading=False, # 延迟加载置为否(立即加载)
|
||
)
|
||
|
||
|
||
@user_context.instructions
|
||
async def query_user_name(ctx: RunContext[Deps]) -> str:
|
||
"""
|
||
查询用户姓名
|
||
"""
|
||
if connection_pool is None:
|
||
return "数据库连接池未初始化"
|
||
|
||
async with connection_pool.acquire() as conn:
|
||
user_name = await conn.fetchval(
|
||
"SELECT user_name FROM bank_customer_service WHERE user_id = $1",
|
||
ctx.deps.user_id,
|
||
)
|
||
if user_name is None:
|
||
return "该用户不存在"
|
||
return f"用户姓名是 {user_name!r}"
|
||
|
||
|
||
@user_context.tool(name="查询账户余额")
|
||
async def query_balance(ctx: RunContext[Deps]) -> str:
|
||
"""
|
||
查询账户余额
|
||
"""
|
||
if connection_pool is None:
|
||
return "数据库连接池未初始化"
|
||
|
||
async with connection_pool.acquire() as conn:
|
||
account_balance = await conn.fetchval(
|
||
"SELECT account_balance FROM bank_customer_service WHERE user_id = $1",
|
||
ctx.deps.user_id,
|
||
)
|
||
if account_balance is None:
|
||
return "该用户不存在"
|
||
return f"该用户的账户余额: {account_balance:.2f}"
|
||
|
||
|
||
# 智能体
|
||
agent = Agent[Deps, str](
|
||
model=DEEPSEEK_V4_FLASH_CHAT_MODEL,
|
||
model_settings=DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS,
|
||
deps_type=Deps,
|
||
output_type=str, # 若输出类型为 str，则智能体输出结果时不进行输出类型校验。若定义输出类型，需考虑处理相应事件
|
||
instructions="你是我方银行的客服智能体，为客户提供客服应答。", # 优先使用 instructions
|
||
capabilities=[user_context],
|
||
)
|
||
|
||
"""
|
||
此为设置输出类型为 Output，智能体输出结果时需校验输出类型，同工具调用：
|
||
PartStartEvent(index=0, part=ToolCallPart(tool_name='final_result', args='', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
FinalResultEvent(tool_name='final_result', tool_call_id='call_3ab0ed20b1114fa0be726723')
|
||
PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='{', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta='"response": " 您好 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta='，刘 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 弼仁 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 先生。', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 您账户 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 的 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 当前余额 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 为 **', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='9.', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta='99 元 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='**。', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 请问还有什么 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 可以帮 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 您的 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 吗？"}', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
PartEndEvent (index=0, part=ToolCallPart (tool_name='final_result', args='{"response": " 您好，刘弼仁先生。您账户的当前余额为 **9.99 元**。请问还有什么可以帮您的吗？"}', tool_call_id='call_3ab0ed20b1114fa0be726723'))
|
||
OutputToolCallEvent (part=ToolCallPart (tool_name='final_result', args='{"response": " 您好，刘弼仁先生。您账户的当前余额为 **9.99 元**。请问还有什么可以帮您的吗？"}', tool_call_id='call_3ab0ed20b1114fa0be726723'), args_valid=True)
|
||
OutputToolResultEvent(part=ToolReturnPart(tool_name='final_result', content='Final result processed.', tool_call_id='call_3ab0ed20b1114fa0be726723', timestamp=datetime.datetime(2026, 9, 17, 10, 34, 0, 799329, tzinfo=datetime.timezone.utc)))
|
||
|
||
此为设置输出类型为 str，智能体输出结果时不进行输出类型校验：
|
||
PartStartEvent(index=0, part=TextPart(content='您好'))
|
||
FinalResultEvent(tool_name=None, tool_call_id=None)
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='，刘'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='弼仁'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='先生。'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='您账户'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='的'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='当前余额'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='为 **'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='9.'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='99'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='元**'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='。请问'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='还有什么'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='可以帮'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='您的吗'))
|
||
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='？'))
|
||
PartEndEvent(index=0, part=TextPart(content='您好，刘弼仁先生。您账户的当前余额为 **9.99元**。请问还有什么可以帮您的吗？'))
|
||
"""
|
||
|
||
|
||
async def run_stream_events(
|
||
user_id: str,
|
||
user_prompt: str | None,
|
||
message_history: list[ModelMessage],
|
||
usage: RunUsage,
|
||
usage_limits: UsageLimits | None = None,
|
||
deferred_tool_results: DeferredToolResults | None = None,
|
||
) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent[str] | None,]:
|
||
"""
|
||
运行并流式输出事件
|
||
"""
|
||
async with agent.run_stream_events(
|
||
user_prompt=user_prompt,
|
||
deps=Deps(user_id=user_id),
|
||
message_history=message_history,
|
||
usage=usage,
|
||
usage_limits=usage_limits,
|
||
deferred_tool_results=deferred_tool_results,
|
||
) as events:
|
||
async for event in events:
|
||
yield event
|