WIN-JHFT4D3SIVT\caoxiaozhu
4251a79062
Introduce a manifest-backed agent registry surface and align graph tests with the new runtime prompt and tool indexing behavior.
382 lines
13 KiB
Python
382 lines
13 KiB
Python
"""
|
|
Jarvis LangGraph Agent 主图定义 - 优化重构版
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
from typing import Literal, Union, List, Any
|
|
|
|
from langchain_core.messages import (
|
|
BaseMessage,
|
|
HumanMessage,
|
|
AIMessage,
|
|
SystemMessage,
|
|
ToolMessage
|
|
)
|
|
from langgraph.graph import StateGraph, END
|
|
|
|
from app.agents.state import AgentState, AgentRole
|
|
from app.agents.prompts import (
|
|
MASTER_SYSTEM_PROMPT,
|
|
SCHEDULE_PLANNER_SYSTEM_PROMPT,
|
|
EXECUTOR_SYSTEM_PROMPT,
|
|
LIBRARIAN_SYSTEM_PROMPT,
|
|
ANALYST_SYSTEM_PROMPT,
|
|
JSON_ACTION_FALLBACK_PROMPT,
|
|
)
|
|
from app.agents.tools import ALL_TOOLS, SUB_COMMANDER_TOOLSETS
|
|
from app.agents.tools.time_reasoning import normalize_tool_time_arguments
|
|
from app.agents.skill_registry import build_skill_context
|
|
from app.services.llm_service import (
|
|
get_llm,
|
|
create_llm_from_config,
|
|
resolve_provider_capabilities,
|
|
default_provider_capabilities
|
|
)
|
|
from app.logging_utils import summarize_llm_config
|
|
|
|
logger = logging.getLogger("jarvis.agent")
|
|
|
|
# ===================== 工具辅助函数 =====================
|
|
|
|
def _get_llm_for_state(state: AgentState):
|
|
"""获取配置好的 LLM 实例"""
|
|
user_llm_config = state.get("user_llm_config")
|
|
llm = create_llm_from_config(user_llm_config) if user_llm_config else get_llm()
|
|
|
|
# 注入解析到的能力
|
|
capabilities = getattr(llm, "_jarvis_provider_capabilities", None)
|
|
if capabilities is None:
|
|
capabilities = resolve_provider_capabilities(user_llm_config) if user_llm_config else default_provider_capabilities()
|
|
|
|
state["provider_capabilities"] = {
|
|
"provider": capabilities.provider,
|
|
"supports_native_tools": capabilities.supports_native_tools,
|
|
"preferred_tool_strategy": capabilities.preferred_tool_strategy,
|
|
}
|
|
return llm, capabilities
|
|
|
|
|
|
def _filter_user_messages(messages: list[BaseMessage]) -> list[BaseMessage]:
|
|
return [m for m in messages if m.type in ("human", "user")]
|
|
|
|
|
|
def _dedupe_tools_by_name(tools: list) -> list:
|
|
deduped_tools = []
|
|
seen_tool_names: set[str] = set()
|
|
for tool in tools:
|
|
if tool.name in seen_tool_names:
|
|
continue
|
|
deduped_tools.append(tool)
|
|
seen_tool_names.add(tool.name)
|
|
return deduped_tools
|
|
|
|
|
|
def _get_role_tools(role: AgentRole) -> list:
|
|
"""获取角色对应的所有可用工具集"""
|
|
if role == AgentRole.SCHEDULE_PLANNER:
|
|
# 合并分析和规划工具
|
|
return _dedupe_tools_by_name(
|
|
SUB_COMMANDER_TOOLSETS["schedule_analysis"]
|
|
+ SUB_COMMANDER_TOOLSETS["schedule_planning"]
|
|
)
|
|
if role == AgentRole.EXECUTOR:
|
|
return _dedupe_tools_by_name(
|
|
SUB_COMMANDER_TOOLSETS["executor_tasks"]
|
|
+ SUB_COMMANDER_TOOLSETS["executor_forum"]
|
|
)
|
|
if role == AgentRole.LIBRARIAN:
|
|
return _dedupe_tools_by_name(
|
|
SUB_COMMANDER_TOOLSETS["librarian_retrieval"]
|
|
+ SUB_COMMANDER_TOOLSETS["librarian_graph"]
|
|
)
|
|
if role == AgentRole.ANALYST:
|
|
return _dedupe_tools_by_name(
|
|
SUB_COMMANDER_TOOLSETS["analyst_progress"]
|
|
+ SUB_COMMANDER_TOOLSETS["analyst_insights"]
|
|
)
|
|
return []
|
|
|
|
|
|
# ===================== 核心执行逻辑 (ReAct) =====================
|
|
|
|
async def call_agent_llm(state: AgentState, role: AgentRole, system_prompt: str) -> dict:
|
|
"""通用的 LLM 调用节点逻辑"""
|
|
llm, capabilities = _get_llm_for_state(state)
|
|
tools = _get_role_tools(role)
|
|
|
|
# 构建消息序列
|
|
messages = []
|
|
|
|
# 1. 系统提示词
|
|
messages.append(SystemMessage(content=system_prompt))
|
|
|
|
# 2. 环境上下文 (时间、记忆等)
|
|
if state.get("current_datetime_context"):
|
|
messages.append(SystemMessage(content=f"当前时间上下文: {state['current_datetime_context']}"))
|
|
|
|
if state.get("memory_context"):
|
|
messages.append(SystemMessage(content=f"长期记忆上下文: {state['memory_context']}"))
|
|
|
|
# 3. 技能增强
|
|
role_skill_key = role.value.replace("agent_", "")
|
|
skill_ctx = build_skill_context(role_skill_key)
|
|
if skill_ctx:
|
|
messages.append(SystemMessage(content=skill_ctx))
|
|
|
|
# 4. 历史对话 (add_messages 已经处理好了)
|
|
messages.extend(state["messages"])
|
|
|
|
# 绑定工具
|
|
if tools and capabilities.supports_native_tools:
|
|
llm_with_tools = llm.bind_tools(tools)
|
|
else:
|
|
llm_with_tools = llm
|
|
if tools: # 如果有工具但不支持原生，注入 JSON Fallback 提示
|
|
messages.append(SystemMessage(content=JSON_ACTION_FALLBACK_PROMPT))
|
|
tool_names = [t.name for t in tools]
|
|
messages.append(SystemMessage(content=f"本次可用工具列表: {', '.join(tool_names)}"))
|
|
|
|
logger.info(
|
|
f"agent_node_started",
|
|
extra={
|
|
"details": {
|
|
"role": role.value,
|
|
"message_count": len(messages),
|
|
"tool_count": len(tools),
|
|
"provider": capabilities.provider
|
|
}
|
|
}
|
|
)
|
|
|
|
# 执行调用
|
|
response = await llm_with_tools.ainvoke(messages)
|
|
|
|
logger.info(
|
|
f"agent_node_finished",
|
|
extra={
|
|
"details": {
|
|
"role": role.value,
|
|
"has_tool_calls": bool(getattr(response, "tool_calls", None)),
|
|
"content_length": len(response.content) if response.content else 0
|
|
}
|
|
}
|
|
)
|
|
|
|
return {"messages": [response]}
|
|
|
|
|
|
async def execute_tools_node(state: AgentState) -> dict:
|
|
"""执行工具调用并返回 ToolMessage 的通用节点"""
|
|
last_message = state["messages"][-1]
|
|
if not hasattr(last_message, "tool_calls") or not last_message.tool_calls:
|
|
return {"messages": []}
|
|
|
|
tool_map = {t.name: t for t in ALL_TOOLS}
|
|
tool_messages = []
|
|
created_entities = []
|
|
|
|
for tool_call in last_message.tool_calls:
|
|
tool_name = tool_call["name"]
|
|
tool_args = tool_call["args"]
|
|
tool_id = tool_call.get("id")
|
|
|
|
logger.info(
|
|
f"tool_execution_started",
|
|
extra={
|
|
"details": {
|
|
"tool_name": tool_name,
|
|
"tool_args": tool_args,
|
|
"tool_id": tool_id
|
|
}
|
|
}
|
|
)
|
|
|
|
try:
|
|
# 时间参数归一化
|
|
normalized_args = normalize_tool_time_arguments(
|
|
tool_name,
|
|
tool_args,
|
|
state.get("current_datetime_context")
|
|
)
|
|
|
|
tool = tool_map.get(tool_name)
|
|
if not tool:
|
|
result = f"Error: Tool {tool_name} not found."
|
|
else:
|
|
result = await tool.ainvoke(normalized_args) if hasattr(tool, "ainvoke") else tool.invoke(normalized_args)
|
|
|
|
# 实体识别（用于业务追踪）
|
|
if any(k in tool_name for k in ["create", "add", "new"]):
|
|
created_entities.append({"tool": tool_name, "result": str(result)})
|
|
|
|
status = "success"
|
|
except Exception as e:
|
|
logger.exception(f"tool_execution_failed: {tool_name}")
|
|
result = f"Error executing tool {tool_name}: {str(e)}"
|
|
status = "failed"
|
|
|
|
tool_messages.append(ToolMessage(
|
|
tool_call_id=tool_id,
|
|
content=str(result),
|
|
name=tool_name
|
|
))
|
|
|
|
logger.info(
|
|
f"tool_execution_finished",
|
|
extra={
|
|
"details": {
|
|
"tool_name": tool_name,
|
|
"status": status,
|
|
"result_preview": str(result)[:200]
|
|
}
|
|
}
|
|
)
|
|
|
|
return {
|
|
"messages": tool_messages,
|
|
"created_entities": state.get("created_entities", []) + created_entities
|
|
}
|
|
|
|
|
|
# ===================== 各角色节点定义 =====================
|
|
|
|
async def master_node(state: AgentState) -> dict:
|
|
"""主控节点：负责意图识别与初步分发"""
|
|
user_messages = _filter_user_messages(state["messages"])
|
|
if not user_messages:
|
|
return {"final_response": "未收到有效输入。"}
|
|
|
|
query = user_messages[-1].content.strip()
|
|
|
|
# 快捷回复逻辑 (保留原有的人性化设计)
|
|
if re.match(r"^(你好|早|在吗|嗨|hi|hello)", query.lower()):
|
|
return {"final_response": "您好。我在。\n\n您把问题给我，我先帮您收束重点，再往下推。", "messages": [AIMessage(content="您好。我在。")]}
|
|
|
|
llm, capabilities = _get_llm_for_state(state)
|
|
|
|
# 路由判断：让 LLM 决定跳转到哪个角色，或者直接回答
|
|
# 这里我们使用一个简洁的提示词让 LLM 输出角色名称或直接回答
|
|
system_msg = SystemMessage(content=MASTER_SYSTEM_PROMPT + "\n\n请直接输出接下来该由哪个 Agent 接手(role_name)，如果直接回答，请正常输出。")
|
|
|
|
response = await llm.ainvoke([system_msg] + state["messages"])
|
|
content = response.content.strip().lower()
|
|
|
|
# 简单的角色映射识别
|
|
roles = {r.value: r for r in AgentRole}
|
|
target_role = None
|
|
for r_val, r_enum in roles.items():
|
|
if r_val in content and len(content) < 50: # 如果内容很短且包含角色名，视为路由
|
|
target_role = r_enum
|
|
break
|
|
|
|
if target_role and target_role != AgentRole.MASTER:
|
|
logger.info(f"master_routing_decided: {target_role.value}")
|
|
return {
|
|
"current_agent": target_role.value,
|
|
"agent_trace": state.get("agent_trace", []) + [target_role.value],
|
|
"messages": [AIMessage(content=f"已分发至 {target_role.value} 处理。")]
|
|
}
|
|
|
|
return {"final_response": response.content, "messages": [response]}
|
|
|
|
|
|
async def planner_node(state: AgentState) -> dict:
|
|
return await call_agent_llm(state, AgentRole.SCHEDULE_PLANNER, SCHEDULE_PLANNER_SYSTEM_PROMPT)
|
|
|
|
async def executor_node(state: AgentState) -> dict:
|
|
return await call_agent_llm(state, AgentRole.EXECUTOR, EXECUTOR_SYSTEM_PROMPT)
|
|
|
|
async def librarian_node(state: AgentState) -> dict:
|
|
return await call_agent_llm(state, AgentRole.LIBRARIAN, LIBRARIAN_SYSTEM_PROMPT)
|
|
|
|
async def analyst_node(state: AgentState) -> dict:
|
|
return await call_agent_llm(state, AgentRole.ANALYST, ANALYST_SYSTEM_PROMPT)
|
|
|
|
|
|
# ===================== 路由逻辑 =====================
|
|
|
|
def route_after_agent(state: AgentState) -> Literal["tools", "__end__"]:
|
|
"""判断 Agent 执行后是该走工具节点还是结束"""
|
|
last_message = state["messages"][-1]
|
|
if hasattr(last_message, "tool_calls") and last_message.tool_calls:
|
|
return "tools"
|
|
return END
|
|
|
|
def route_master(state: AgentState) -> str:
|
|
"""主控路由逻辑"""
|
|
if state.get("final_response"):
|
|
return END
|
|
return state.get("current_agent", END)
|
|
|
|
|
|
# ===================== 图构建 =====================
|
|
|
|
def create_agent_graph(callbacks: list | None = None):
|
|
workflow = StateGraph(AgentState)
|
|
|
|
# 添加节点
|
|
workflow.add_node(AgentRole.MASTER.value, master_node)
|
|
workflow.add_node(AgentRole.SCHEDULE_PLANNER.value, planner_node)
|
|
workflow.add_node(AgentRole.EXECUTOR.value, executor_node)
|
|
workflow.add_node(AgentRole.LIBRARIAN.value, librarian_node)
|
|
workflow.add_node(AgentRole.ANALYST.value, analyst_node)
|
|
workflow.add_node("tools", execute_tools_node)
|
|
|
|
# 设置入口
|
|
workflow.set_entry_point(AgentRole.MASTER.value)
|
|
|
|
# 主控分发逻辑
|
|
workflow.add_conditional_edges(
|
|
AgentRole.MASTER.value,
|
|
route_master,
|
|
{
|
|
AgentRole.SCHEDULE_PLANNER.value: AgentRole.SCHEDULE_PLANNER.value,
|
|
AgentRole.EXECUTOR.value: AgentRole.EXECUTOR.value,
|
|
AgentRole.LIBRARIAN.value: AgentRole.LIBRARIAN.value,
|
|
AgentRole.ANALYST.value: AgentRole.ANALYST.value,
|
|
END: END
|
|
}
|
|
)
|
|
|
|
# 各角色节点的 ReAct 循环
|
|
for role in [AgentRole.SCHEDULE_PLANNER, AgentRole.EXECUTOR, AgentRole.LIBRARIAN, AgentRole.ANALYST]:
|
|
workflow.add_conditional_edges(
|
|
role.value,
|
|
route_after_agent,
|
|
{
|
|
"tools": "tools",
|
|
END: END
|
|
}
|
|
)
|
|
|
|
# 工具执行完后回到当前 Agent 角色继续处理
|
|
workflow.add_conditional_edges(
|
|
"tools",
|
|
lambda s: s.get("current_agent", AgentRole.MASTER.value),
|
|
{
|
|
AgentRole.SCHEDULE_PLANNER.value: AgentRole.SCHEDULE_PLANNER.value,
|
|
AgentRole.EXECUTOR.value: AgentRole.EXECUTOR.value,
|
|
AgentRole.LIBRARIAN.value: AgentRole.LIBRARIAN.value,
|
|
AgentRole.ANALYST.value: AgentRole.ANALYST.value,
|
|
}
|
|
)
|
|
|
|
# 编译
|
|
if callbacks:
|
|
return workflow.compile(callbacks=callbacks)
|
|
return workflow.compile()
|
|
|
|
|
|
_agent_graph = None
|
|
|
|
def get_agent_graph(callbacks: list | None = None):
|
|
global _agent_graph
|
|
if _agent_graph is None:
|
|
from app.config_tracing import get_langsmith_callbacks
|
|
langsmith_callbacks = get_langsmith_callbacks()
|
|
all_callbacks = (callbacks or []) + langsmith_callbacks
|
|
_agent_graph = create_agent_graph(callbacks=all_callbacks or None)
|
|
return _agent_graph
|