zhanghongyu_0626
3995cb44d8
- Added `auth.py` for mock login and JWT issuance. - Introduced `chat.py` for handling chat requests with role-based access control. - Enhanced `main.py` to include new routers and middleware for tracing. - Implemented input validation in `input_guard.py` to prevent SQL injection. - Created repositories for managing agent sessions and audit logs. - Added exception handling for authorization errors. - Updated settings to include JWT configuration. - Introduced tests for authentication and input validation.
65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
"""Agent 编排：LangGraph StateGraph 最小骨架。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import TypedDict
|
||
|
||
from langgraph.graph import END, StateGraph
|
||
|
||
from app.model.schemas import AgentType, AuthContext
|
||
|
||
DISCLAIMER = (
|
||
"本内容仅为投资分析参考，不构成任何直接投资建议，不构成对任何产品的收益承诺，"
|
||
"据此操作风险自负，请谨慎对待。"
|
||
)
|
||
|
||
|
||
class AgentState(TypedDict):
|
||
message: str
|
||
reply: str
|
||
agent_type: AgentType
|
||
actor_id: str
|
||
customer_id: str | None
|
||
|
||
|
||
def _build_graph():
|
||
graph = StateGraph(AgentState)
|
||
|
||
def respond(state: AgentState) -> AgentState:
|
||
agent = state["agent_type"]
|
||
prefix = {
|
||
"customer": "客户财富助手",
|
||
"advisor": "代理人助手",
|
||
"analyst": "数据分析助手",
|
||
"risk": "风控监测助手",
|
||
}.get(agent, "助手")
|
||
target = f"（客户 {state['customer_id']}）" if state.get("customer_id") else ""
|
||
reply = f"【{prefix}】{target}已收到：{state['message']}"
|
||
if agent in ("advisor", "analyst") and any(
|
||
k in state["message"] for k in ("报告", "建议", "推荐")
|
||
):
|
||
reply = f"{reply}\n\n{DISCLAIMER}"
|
||
return {"reply": reply}
|
||
|
||
graph.add_node("respond", respond)
|
||
graph.set_entry_point("respond")
|
||
graph.add_edge("respond", END)
|
||
return graph.compile()
|
||
|
||
|
||
_GRAPH = _build_graph()
|
||
|
||
|
||
def run_chat(ctx: AuthContext, message: str, customer_id: str | None) -> tuple[str, bool]:
|
||
state: AgentState = {
|
||
"message": message,
|
||
"reply": "",
|
||
"agent_type": ctx.agent_type,
|
||
"actor_id": ctx.sub,
|
||
"customer_id": customer_id,
|
||
}
|
||
result = _GRAPH.invoke(state)
|
||
reply = result["reply"]
|
||
has_disclaimer = DISCLAIMER in reply
|
||
return reply, has_disclaimer
|