471 lines
18 KiB
Python
471 lines
18 KiB
Python
from types import SimpleNamespace
|
||
|
||
from langchain_core.messages import AIMessage, HumanMessage
|
||
|
||
from app.agents.graph import (
|
||
_choose_sub_commander,
|
||
_parse_json_action,
|
||
_route_agent_from_user_query,
|
||
_run_sub_commander,
|
||
master_node,
|
||
)
|
||
from app.agents.tools.time_reasoning import resolve_time_expression
|
||
from app.agents.state import AgentRole
|
||
|
||
|
||
|
||
|
||
def _base_state(message: str, user_llm_config: dict | None = None) -> dict:
|
||
return {
|
||
'messages': [HumanMessage(content=message)],
|
||
'user_id': 'u1',
|
||
'conversation_id': 'c1',
|
||
'current_agent': AgentRole.MASTER,
|
||
'active_agents': [AgentRole.MASTER],
|
||
'current_sub_commander': None,
|
||
'active_sub_commanders': [],
|
||
'sub_commander_trace': [],
|
||
'pending_tasks': [],
|
||
'completed_tasks': [],
|
||
'tool_calls': [],
|
||
'last_tool_result': None,
|
||
'action_results': [],
|
||
'created_entities': [],
|
||
'tool_strategy_used': None,
|
||
'provider_capabilities': None,
|
||
'fallback_parse_error': None,
|
||
'knowledge_context': None,
|
||
'graph_context': None,
|
||
'schedule_context_summary': None,
|
||
'plan': None,
|
||
'plan_steps': [],
|
||
'analysis_report': None,
|
||
'final_response': None,
|
||
'should_respond': True,
|
||
'memory_context': None,
|
||
'current_datetime_context': 'CURRENT_TIME: 2026-03-28T12:00:00+08:00',
|
||
'current_datetime_reference': {'current_time_iso': '2026-03-28T12:00:00+08:00', 'current_date_iso': '2026-03-28', 'timezone': 'UTC'},
|
||
'user_llm_config': user_llm_config,
|
||
}
|
||
|
||
|
||
class FakeFallbackLLM:
|
||
def __init__(self, first_content: str, followup_content: str = '已创建提醒：开会，时间为 2026-03-29 09:00（按当前时间理解为“明天早上9点”）。'):
|
||
self.first_content = first_content
|
||
self.followup_content = followup_content
|
||
self.calls = 0
|
||
|
||
async def ainvoke(self, messages):
|
||
self.calls += 1
|
||
if self.calls == 1:
|
||
return AIMessage(content=self.first_content)
|
||
return AIMessage(content=self.followup_content)
|
||
|
||
def bind_tools(self, tools):
|
||
raise AssertionError('bind_tools should not be called in JSON fallback mode')
|
||
|
||
|
||
class FakeNativeBoundLLM:
|
||
async def ainvoke(self, messages):
|
||
return AIMessage(
|
||
content='',
|
||
tool_calls=[
|
||
{
|
||
'id': 'call_1',
|
||
'name': 'create_reminder',
|
||
'args': {'title': '开会', 'reminder_at': '明天 09:00'},
|
||
}
|
||
],
|
||
)
|
||
|
||
|
||
class FakeNativeLLM:
|
||
def __init__(self):
|
||
self.bound = FakeNativeBoundLLM()
|
||
self.tool_binding_count = 0
|
||
self.calls = 0
|
||
self._jarvis_provider_capabilities = SimpleNamespace(provider='openai', supports_native_tools=True, preferred_tool_strategy='native')
|
||
|
||
def bind_tools(self, tools):
|
||
self.tool_binding_count += 1
|
||
return self.bound
|
||
|
||
async def ainvoke(self, messages):
|
||
self.calls += 1
|
||
return AIMessage(content='已创建提醒：开会，时间为 2026-03-29 09:00（按当前时间理解为“明天早上9点”）。')
|
||
|
||
|
||
class FakeTool:
|
||
def __init__(self, name: str, result: str):
|
||
self.name = name
|
||
self.result = result
|
||
self.invocations: list[dict] = []
|
||
|
||
def invoke(self, args: dict):
|
||
self.invocations.append(args)
|
||
return self.result
|
||
|
||
|
||
class CapturingLLM:
|
||
def __init__(self, content: str = '{"mode":"final","final_response":"好的。"}'):
|
||
self.content = content
|
||
self.messages = None
|
||
self._jarvis_provider_capabilities = SimpleNamespace(provider='ollama', supports_native_tools=False, preferred_tool_strategy='json_fallback')
|
||
|
||
async def ainvoke(self, messages):
|
||
self.messages = messages
|
||
return AIMessage(content=self.content)
|
||
|
||
|
||
class FailIfCalledLLM:
|
||
async def ainvoke(self, messages):
|
||
raise AssertionError('LLM should not be called for simple greetings')
|
||
|
||
|
||
async def test_master_node_returns_stable_reply_for_simple_greeting(monkeypatch):
|
||
monkeypatch.setattr('app.agents.graph._get_llm_for_state', lambda state: FailIfCalledLLM())
|
||
|
||
state = {
|
||
'messages': [HumanMessage(content='你好')],
|
||
'user_id': 'u1',
|
||
'conversation_id': 'c1',
|
||
'current_agent': AgentRole.MASTER,
|
||
'active_agents': [AgentRole.MASTER],
|
||
'pending_tasks': [],
|
||
'completed_tasks': [],
|
||
'tool_calls': [],
|
||
'last_tool_result': None,
|
||
'knowledge_context': None,
|
||
'graph_context': None,
|
||
'plan': None,
|
||
'plan_steps': [],
|
||
'analysis_report': None,
|
||
'final_response': None,
|
||
'should_respond': True,
|
||
'memory_context': None,
|
||
'user_llm_config': None,
|
||
}
|
||
|
||
result = await master_node(state)
|
||
|
||
assert result['final_response'] == '您好。我在。\n\n您把问题给我，我先帮您收束重点，再往下推。'
|
||
assert result['current_agent'] == AgentRole.MASTER
|
||
assert result['active_agents'] == [AgentRole.MASTER]
|
||
|
||
|
||
async def test_master_node_returns_stable_reply_for_identity_question(monkeypatch):
|
||
monkeypatch.setattr('app.agents.graph._get_llm_for_state', lambda state: FailIfCalledLLM())
|
||
|
||
state = {
|
||
'messages': [HumanMessage(content='你是谁')],
|
||
'user_id': 'u1',
|
||
'conversation_id': 'c1',
|
||
'current_agent': AgentRole.MASTER,
|
||
'active_agents': [AgentRole.MASTER],
|
||
'pending_tasks': [],
|
||
'completed_tasks': [],
|
||
'tool_calls': [],
|
||
'last_tool_result': None,
|
||
'knowledge_context': None,
|
||
'graph_context': None,
|
||
'plan': None,
|
||
'plan_steps': [],
|
||
'analysis_report': None,
|
||
'final_response': None,
|
||
'should_respond': True,
|
||
'memory_context': None,
|
||
'user_llm_config': None,
|
||
}
|
||
|
||
result = await master_node(state)
|
||
|
||
assert result['final_response'] == '我是 Jarvis。\n\n比起做一个泛泛的助手，我更像您的判断型协作伙伴：帮您看清问题、压缩路径、把事情往前推进。'
|
||
assert result['current_agent'] == AgentRole.MASTER
|
||
assert result['active_agents'] == [AgentRole.MASTER]
|
||
|
||
|
||
async def test_master_node_returns_stable_reply_for_identity_question_with_punctuation(monkeypatch):
|
||
monkeypatch.setattr('app.agents.graph._get_llm_for_state', lambda state: FailIfCalledLLM())
|
||
|
||
state = {
|
||
'messages': [HumanMessage(content='你是谁？')],
|
||
'user_id': 'u1',
|
||
'conversation_id': 'c1',
|
||
'current_agent': AgentRole.MASTER,
|
||
'active_agents': [AgentRole.MASTER],
|
||
'pending_tasks': [],
|
||
'completed_tasks': [],
|
||
'tool_calls': [],
|
||
'last_tool_result': None,
|
||
'knowledge_context': None,
|
||
'graph_context': None,
|
||
'plan': None,
|
||
'plan_steps': [],
|
||
'analysis_report': None,
|
||
'final_response': None,
|
||
'should_respond': True,
|
||
'memory_context': None,
|
||
'user_llm_config': None,
|
||
}
|
||
|
||
result = await master_node(state)
|
||
|
||
assert result['final_response'] == '我是 Jarvis。\n\n比起做一个泛泛的助手，我更像您的判断型协作伙伴：帮您看清问题、压缩路径、把事情往前推进。'
|
||
assert result['current_agent'] == AgentRole.MASTER
|
||
assert result['active_agents'] == [AgentRole.MASTER]
|
||
|
||
|
||
async def test_master_node_returns_stable_reply_for_identity_question_with_particle(monkeypatch):
|
||
monkeypatch.setattr('app.agents.graph._get_llm_for_state', lambda state: FailIfCalledLLM())
|
||
|
||
state = {
|
||
'messages': [HumanMessage(content='你是谁啊')],
|
||
'user_id': 'u1',
|
||
'conversation_id': 'c1',
|
||
'current_agent': AgentRole.MASTER,
|
||
'active_agents': [AgentRole.MASTER],
|
||
'pending_tasks': [],
|
||
'completed_tasks': [],
|
||
'tool_calls': [],
|
||
'last_tool_result': None,
|
||
'knowledge_context': None,
|
||
'graph_context': None,
|
||
'plan': None,
|
||
'plan_steps': [],
|
||
'analysis_report': None,
|
||
'final_response': None,
|
||
'should_respond': True,
|
||
'memory_context': None,
|
||
'user_llm_config': None,
|
||
}
|
||
|
||
result = await master_node(state)
|
||
|
||
assert result['final_response'] == '我是 Jarvis。\n\n比起做一个泛泛的助手，我更像您的判断型协作伙伴：帮您看清问题、压缩路径、把事情往前推进。'
|
||
assert result['current_agent'] == AgentRole.MASTER
|
||
assert result['active_agents'] == [AgentRole.MASTER]
|
||
|
||
|
||
async def test_master_node_returns_stable_reply_for_capability_question(monkeypatch):
|
||
monkeypatch.setattr('app.agents.graph._get_llm_for_state', lambda state: FailIfCalledLLM())
|
||
|
||
state = {
|
||
'messages': [HumanMessage(content='你能做什么')],
|
||
'user_id': 'u1',
|
||
'conversation_id': 'c1',
|
||
'current_agent': AgentRole.MASTER,
|
||
'active_agents': [AgentRole.MASTER],
|
||
'pending_tasks': [],
|
||
'completed_tasks': [],
|
||
'tool_calls': [],
|
||
'last_tool_result': None,
|
||
'knowledge_context': None,
|
||
'graph_context': None,
|
||
'plan': None,
|
||
'plan_steps': [],
|
||
'analysis_report': None,
|
||
'final_response': None,
|
||
'should_respond': True,
|
||
'memory_context': None,
|
||
'user_llm_config': None,
|
||
}
|
||
|
||
result = await master_node(state)
|
||
|
||
assert result['final_response'] == '主要做三件事。\n- 帮您判断：看问题本质、梳理取舍、给出方向\n- 帮您收束：把复杂内容理顺，把重点拎出来\n- 帮您推进：拆任务、定步骤、把下一步变清楚\n\n如果您现在有具体目标，我可以直接进入处理。'
|
||
assert result['current_agent'] == AgentRole.MASTER
|
||
assert result['active_agents'] == [AgentRole.MASTER]
|
||
|
||
|
||
def test_choose_sub_commander_routes_schedule_requests_to_schedule_planning():
|
||
assert _choose_sub_commander(AgentRole.SCHEDULE_PLANNER, '帮我安排一下这周计划') == 'schedule_planning'
|
||
|
||
|
||
def test_choose_sub_commander_routes_focus_requests_to_schedule_analysis():
|
||
assert _choose_sub_commander(AgentRole.SCHEDULE_PLANNER, '基于最近对话帮我判断该聚焦什么') == 'schedule_analysis'
|
||
|
||
|
||
def test_route_agent_from_user_query_routes_knowledge_requests_to_librarian():
|
||
assert _route_agent_from_user_query('帮我搜索知识库里的项目资料') == AgentRole.LIBRARIAN
|
||
|
||
|
||
def test_route_agent_from_user_query_routes_schedule_requests_to_schedule_planner():
|
||
assert _route_agent_from_user_query('明天提醒我开会') == AgentRole.SCHEDULE_PLANNER
|
||
|
||
|
||
def test_route_agent_from_user_query_routes_explicit_month_day_milestone_to_schedule_planner():
|
||
assert _route_agent_from_user_query('3月29日，对话系统交付节点') == AgentRole.SCHEDULE_PLANNER
|
||
|
||
|
||
def test_choose_sub_commander_routes_explicit_month_day_milestone_to_schedule_planning():
|
||
assert _choose_sub_commander(AgentRole.SCHEDULE_PLANNER, '3月29日，对话系统交付节点') == 'schedule_planning'
|
||
|
||
|
||
|
||
|
||
def test_parse_json_action_extracts_tool_calls_from_fenced_json():
|
||
parsed = _parse_json_action(
|
||
'```json\n{"mode":"tool_call","tool_calls":[{"name":"create_reminder","arguments":{"title":"开会","reminder_at":"明天 09:00"}}]}\n```',
|
||
['create_reminder'],
|
||
)
|
||
|
||
assert parsed == {
|
||
'mode': 'tool_call',
|
||
'tool_calls': [
|
||
{
|
||
'name': 'create_reminder',
|
||
'args': {'title': '开会', 'reminder_at': '明天 09:00'},
|
||
'reason': None,
|
||
}
|
||
],
|
||
}
|
||
|
||
|
||
def test_parse_json_action_returns_none_for_invalid_or_unknown_payload():
|
||
assert _parse_json_action('not json', ['create_reminder']) is None
|
||
assert _parse_json_action('{"mode":"tool_call","tool_calls":[{"name":"unknown","arguments":{}}]}', ['create_reminder']) is None
|
||
|
||
|
||
def test_parse_json_action_tolerates_prefix_and_suffix_text():
|
||
parsed = _parse_json_action(
|
||
'好的，下面是 JSON：\n```json\n{"mode":"tool_call","tool_calls":[{"name":"create_reminder","arguments":{"title":"开会","reminder_at":"明天 09:00"}}]}\n```\n谢谢',
|
||
['create_reminder'],
|
||
)
|
||
assert parsed is not None
|
||
assert parsed['mode'] == 'tool_call'
|
||
assert parsed['tool_calls'][0]['name'] == 'create_reminder'
|
||
|
||
|
||
def test_parse_json_action_accepts_parameters_alias_for_tool_calls():
|
||
parsed = _parse_json_action(
|
||
'{"mode":"tool_call","tool_calls":[{"name":"create_reminder","parameters":{"title":"收被子","reminder_at":"2026-03-29T09:00:00+08:00"}}]}',
|
||
['create_reminder'],
|
||
)
|
||
|
||
assert parsed == {
|
||
'mode': 'tool_call',
|
||
'tool_calls': [
|
||
{
|
||
'name': 'create_reminder',
|
||
'args': {'title': '收被子', 'reminder_at': '2026-03-29T09:00:00+08:00'},
|
||
'reason': None,
|
||
}
|
||
],
|
||
}
|
||
|
||
|
||
async def test_run_sub_commander_uses_json_fallback_for_non_native_provider(monkeypatch):
|
||
fake_llm = FakeFallbackLLM(
|
||
'{"mode":"tool_call","tool_calls":[{"name":"create_reminder","arguments":{"title":"开会","reminder_at":"明天 09:00"}}]}'
|
||
)
|
||
fake_tool = FakeTool('create_reminder', '成功创建 reminder: 开会 @ 明天 09:00')
|
||
|
||
monkeypatch.setattr('app.agents.graph._get_llm_for_state', lambda state: fake_llm)
|
||
monkeypatch.setitem(
|
||
__import__('app.agents.graph', fromlist=['SUB_COMMANDER_TOOLSETS']).SUB_COMMANDER_TOOLSETS,
|
||
'schedule_planning',
|
||
[fake_tool],
|
||
)
|
||
|
||
state = _base_state('明天 9 点提醒我开会', {'provider': 'ollama', 'model': 'qwen2.5'})
|
||
state['current_agent'] = AgentRole.SCHEDULE_PLANNER
|
||
|
||
result = await _run_sub_commander(
|
||
state,
|
||
AgentRole.SCHEDULE_PLANNER,
|
||
'manager prompt',
|
||
'明天 9 点提醒我开会',
|
||
use_tools=True,
|
||
)
|
||
|
||
assert result['tool_strategy_used'] == 'json_fallback'
|
||
assert fake_tool.invocations == [{'title': '开会', 'reminder_at': '2026-03-29T09:00:00'}]
|
||
assert result['tool_calls'][0]['name'] == 'create_reminder'
|
||
assert result['created_entities'][0]['type'] == 'reminder'
|
||
assert result['fallback_parse_error'] is None
|
||
assert result['final_response'] == '已创建提醒：开会，时间为 2026-03-29 09:00（按当前时间理解为“明天早上9点”）。'
|
||
|
||
|
||
async def test_run_sub_commander_includes_current_datetime_context_in_system_messages(monkeypatch):
|
||
fake_llm = CapturingLLM('{"mode":"final","final_response":"好的。"}')
|
||
monkeypatch.setattr('app.agents.graph._get_llm_for_state', lambda state: fake_llm)
|
||
|
||
state = _base_state('明天 9 点提醒我开会', {'provider': 'ollama', 'model': 'qwen2.5'})
|
||
state['current_agent'] = AgentRole.SCHEDULE_PLANNER
|
||
state['current_datetime_context'] = 'CURRENT_TIME: 2026-03-28T12:00:00+08:00'
|
||
|
||
await _run_sub_commander(
|
||
state,
|
||
AgentRole.SCHEDULE_PLANNER,
|
||
'manager prompt',
|
||
'明天 9 点提醒我开会',
|
||
use_tools=True,
|
||
)
|
||
|
||
assert fake_llm.messages is not None
|
||
assert any(
|
||
getattr(m, 'type', None) == 'system' and 'CURRENT_TIME:' in str(getattr(m, 'content', ''))
|
||
for m in fake_llm.messages
|
||
)
|
||
|
||
|
||
async def test_run_sub_commander_uses_web_search_in_json_fallback(monkeypatch):
|
||
fake_llm = FakeFallbackLLM(
|
||
'{"mode":"tool_call","tool_calls":[{"name":"web_search","arguments":{"query":"Jarvis 最新模型更新","top_k":2}}]}',
|
||
'我查了外部网页，下面是最新结果摘要。',
|
||
)
|
||
fake_tool = FakeTool('web_search', '成功搜索到 2 条网页结果')
|
||
|
||
monkeypatch.setattr('app.agents.graph._get_llm_for_state', lambda state: fake_llm)
|
||
monkeypatch.setitem(
|
||
__import__('app.agents.graph', fromlist=['SUB_COMMANDER_TOOLSETS']).SUB_COMMANDER_TOOLSETS,
|
||
'librarian_retrieval',
|
||
[fake_tool],
|
||
)
|
||
|
||
state = _base_state('帮我上网查一下 Jarvis 最新模型更新', {'provider': 'ollama', 'model': 'qwen2.5'})
|
||
state['current_agent'] = AgentRole.LIBRARIAN
|
||
|
||
result = await _run_sub_commander(
|
||
state,
|
||
AgentRole.LIBRARIAN,
|
||
'manager prompt',
|
||
'帮我上网查一下 Jarvis 最新模型更新',
|
||
use_tools=True,
|
||
summary_target='knowledge_context',
|
||
)
|
||
|
||
assert result['tool_strategy_used'] == 'json_fallback'
|
||
assert fake_tool.invocations == [{'query': 'Jarvis 最新模型更新', 'top_k': 2}]
|
||
assert result['tool_calls'][0]['name'] == 'web_search'
|
||
assert result['last_tool_result'] == '[web_search] 成功搜索到 2 条网页结果'
|
||
assert result['final_response'] == '我查了外部网页，下面是最新结果摘要。'
|
||
|
||
|
||
fake_llm = FakeNativeLLM()
|
||
fake_tool = FakeTool('create_reminder', '成功创建 reminder: 开会 @ 明天 09:00')
|
||
|
||
monkeypatch.setattr('app.agents.graph._get_llm_for_state', lambda state: fake_llm)
|
||
monkeypatch.setitem(
|
||
__import__('app.agents.graph', fromlist=['SUB_COMMANDER_TOOLSETS']).SUB_COMMANDER_TOOLSETS,
|
||
'schedule_planning',
|
||
[fake_tool],
|
||
)
|
||
|
||
state = _base_state('明天 9 点提醒我开会', {'provider': 'openai', 'model': 'gpt-4o'})
|
||
state['current_agent'] = AgentRole.SCHEDULE_PLANNER
|
||
|
||
result = await _run_sub_commander(
|
||
state,
|
||
AgentRole.SCHEDULE_PLANNER,
|
||
'manager prompt',
|
||
'明天 9 点提醒我开会',
|
||
use_tools=True,
|
||
)
|
||
|
||
assert result['tool_strategy_used'] == 'native'
|
||
assert fake_llm.tool_binding_count == 1
|
||
assert fake_tool.invocations == [{'title': '开会', 'reminder_at': '2026-03-29T09:00:00'}]
|
||
assert result['created_entities'][0]['type'] == 'reminder'
|
||
assert result['final_response'] == '已创建提醒：开会，时间为 2026-03-29 09:00（按当前时间理解为“明天早上9点”）。'
|