1

agent/application/application.py

@@ -3,7 +3,10 @@
 应用主入口
 """
 import reflex as rx
+
 from application.pages.index import index
+from application.workshop.database import connection_pool_init
+
 
 app = rx.App(
     theme=rx.theme(color_mode="light"),
@@ -36,5 +39,8 @@ app = rx.App(
         },  # 覆盖占位符样式
     },  # 自定义主题颜色
 )  # 此处变量名需使用 app ，具体原因目前尚不清楚
-# 注册首页路由
+
+# 添加页面：首页
 app.add_page(component=index(), route="/")
+# 注册声明周期任务：连接池初始化
+app.register_lifespan_task(connection_pool_init)

agent/application/pages/conversation.py

@@ -226,6 +226,24 @@ def conversation_history(
     )
 
 
+def workflow_button(workflow_type: WorkflowType):
+    """
+    工作流按钮
+    :param workflow_type: 工作类型
+    :return: rx.Component
+    """
+    return rx.box(
+        workflow_type.value,
+        # 点击事件：设置工作流
+        on_click=lambda: ConversationState.set_workflow(workflow_type),
+        padding="10px 16px",
+        background_color="var(--prismui-background-color-6)",
+        border_radius="var(--prismui-border-radius-9)",
+        color="var(--prismui-color-10)",
+        cursor="pointer",
+    )
+
+
 def use_guidance() -> rx.Component:
     """
     使用指引
@@ -276,18 +294,10 @@ def use_guidance() -> rx.Component:
                 font_weight="var(--prismui-font-weight-3)",
             ),
             rx.hstack(
-                rx.box(
-                    "预定航班",
-                    # 点击事件：设置工作流为预定航班
-                    on_click=lambda: ConversationState.set_workflow(
-                        WorkflowType.BOOK_FLIGHT
-                    ),
-                    padding="10px 16px",
-                    background_color="var(--prismui-background-color-6)",
-                    border_radius="var(--prismui-border-radius-9)",
-                    color="var(--prismui-color-10)",
-                    cursor="pointer",
-                ),
+                # 预定航班按钮
+                workflow_button(WorkflowType.BOOK_FLIGHT),
+                # 银行客服按钮
+                workflow_button(WorkflowType.BANK_CUSTOMER_SERVICE),
                 flex_wrap="wrap",
                 align_items="center",
                 gap="12px",

agent/application/states/conversation.py

@@ -357,6 +357,8 @@ class ConversationState(rx.State):
                     message_history=message_history,
                     usage=usage,
                 )
+
+            # 预定航班
             case WorkflowType.BOOK_FLIGHT:
                 from application.workshop.book_flight import (
                     run_stream_events,
@@ -371,6 +373,30 @@ class ConversationState(rx.State):
                     deferred_tool_results=deferred_tool_results,
                 )
 
+            # 银行客户服务
+            case WorkflowType.BANK_CUSTOMER_SERVICE:
+                from application.workshop.bank_customer_service import (
+                    run_stream_events,
+                    Deps,
+                )
+                from application.workshop.database import connection_pool
+
+                # 若连接池为空则直接返回
+                if not connection_pool:
+                    return
+
+                deps = Deps(
+                    account_id=conversation.workflow.deps["account_id"],
+                    connection_pool=connection_pool,
+                )
+                # 运行并流式输出事件
+                stream_events = run_stream_events(
+                    deps=deps,
+                    user_prompt=user_prompt,
+                    message_history=message_history,
+                    usage=usage,
+                )
+
         # 消息列表
         messages: list[Message] = []
         # 获取运行流式输出事件
@@ -540,8 +566,8 @@ class ConversationState(rx.State):
                             messages.append(message)
 
                     else:
-                        # 若工作流类型非对话则设置为对话
-                        if conversation.workflow.type != WorkflowType.TALK:
+                        # 若工作流类型为预定航班则设置为对话
+                        if conversation.workflow.type == WorkflowType.BOOK_FLIGHT:
                             conversation.workflow = Workflow()
                         # 将正在运行设置为否
                         conversation.is_running = False
@@ -634,6 +660,7 @@ class ConversationState(rx.State):
             return
 
         match type:
+            # 预定航班
             case WorkflowType.BOOK_FLIGHT:
                 from application.workshop.book_flight import deps
 
@@ -644,6 +671,13 @@ class ConversationState(rx.State):
                     f"帮我查询并预定在 {workflow.deps.date.strftime('%Y-%m-%d')} 从 {workflow.deps.origin_airport_code} 到 {workflow.deps.destination_airport_code} 的航班"
                 )
 
+            # 银行客户服务
+            case WorkflowType.BANK_CUSTOMER_SERVICE:
+                # 初始化银行客户服务工作流
+                workflow = Workflow(type=type, deps={"account_id": self.user_id})
+                # 设置用户提示词
+                self.set_user_prompt(f"帮我查询账户余额")
+
             case _:
                 # 初始化对话工作流
                 workflow = Workflow()

agent/application/states/models.py

@@ -58,8 +58,9 @@ class WorkflowType(StrEnum):
     工作类型枚举
     """
 
-    TALK = "talk"
+    TALK = "对话"
     BOOK_FLIGHT = "预定航班"
+    BANK_CUSTOMER_SERVICE = "银行客户服务"
 
 
 class Workflow(BaseModel):

agent/application/workshop/bank_customer_service.py

@@ -0,0 +1,149 @@
+# -*- 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

agent/application/workshop/book_flight.py

@@ -8,7 +8,6 @@ from typing import Any, AsyncGenerator, cast
 from pydantic import BaseModel, Field, field_validator
 from pydantic_ai import (
     Agent,
-    AgentRunResult,
     ApprovalRequired,
     DeferredToolRequests,
     DeferredToolResults,
@@ -25,7 +24,7 @@ from pydantic_ai.messages import (
     ToolCallPart,
     ToolReturnPart,
 )
-from pydantic_ai.run import AgentRunResultEvent
+from pydantic_ai.run import AgentRunResultEvent, AgentRunResult
 
 from application.workshop.events import AgentRuntimeEvent
 from application.workshop.models import (
@@ -348,7 +347,7 @@ async def run_stream_events(
                                         tool_call_id=tool_call_id,
                                     )
                                     yield agent_runtime_event
-                            
+
                             yield event  # 智能体运行时依赖 event，原则上不可修改
 
                 case AgentRunResultEvent(result=result):

agent/application/workshop/database.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+"""
+数据库模块
+"""
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+from asyncpg import Pool, create_pool
+
+# 连接池
+connection_pool: Pool | None = None
+
+
+@asynccontextmanager
+async def connection_pool_init() -> AsyncIterator[None]:
+    """
+    连接池初始化
+    """
+    global connection_pool
+
+    # 创建连接池
+    connection_pool = await create_pool(
+        "postgres://root:PG198752@liubiren.cloud:5201/database",
+        min_size=1,
+        max_size=5,
+        server_settings={"statement_timeout": "5000"},
+    )
+    yield
+    if connection_pool:
+        await connection_pool.close()

agent/application/workshop/events.py

@@ -13,4 +13,6 @@ class AgentRuntimeEvent(CustomEvent):
     智能体运行时事件
     """
 
+    # tool_call_id 无需自定义，继承 CustomEvent 中 tool_call_id
+
     content: str = ""