diff --git a/agent/application/application.py b/agent/application/application.py index 8175713..afe4427 100644 --- a/agent/application/application.py +++ b/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) diff --git a/agent/application/pages/conversation.py b/agent/application/pages/conversation.py index a19e48a..3803022 100644 --- a/agent/application/pages/conversation.py +++ b/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", diff --git a/agent/application/states/conversation.py b/agent/application/states/conversation.py index f6411bc..7b96929 100644 --- a/agent/application/states/conversation.py +++ b/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() diff --git a/agent/application/states/models.py b/agent/application/states/models.py index bd8857b..0a04d1f 100644 --- a/agent/application/states/models.py +++ b/agent/application/states/models.py @@ -58,8 +58,9 @@ class WorkflowType(StrEnum): 工作类型枚举 """ - TALK = "talk" + TALK = "对话" BOOK_FLIGHT = "预定航班" + BANK_CUSTOMER_SERVICE = "银行客户服务" class Workflow(BaseModel): diff --git a/agent/application/workshop/bank_customer_service.py b/agent/application/workshop/bank_customer_service.py new file mode 100644 index 0000000..a4dc955 --- /dev/null +++ b/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 diff --git a/agent/application/workshop/book_flight.py b/agent/application/workshop/book_flight.py index cc1c57a..d13e4a2 100644 --- a/agent/application/workshop/book_flight.py +++ b/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): diff --git a/agent/application/workshop/database.py b/agent/application/workshop/database.py new file mode 100644 index 0000000..867f5e2 --- /dev/null +++ b/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() diff --git a/agent/application/workshop/events.py b/agent/application/workshop/events.py index e3b5e57..5125dd2 100644 --- a/agent/application/workshop/events.py +++ b/agent/application/workshop/events.py @@ -13,4 +13,6 @@ class AgentRuntimeEvent(CustomEvent): 智能体运行时事件 """ + # tool_call_id 无需自定义，继承 CustomEvent 中 tool_call_id + content: str = "" diff --git a/agent/database.db b/agent/database.db index 18d31a9..c83cc25 100644 Binary files a/agent/database.db and b/agent/database.db differ