from langchain_core.messages import HumanMessage from agent.agents.conversation import ConversationAgent from agent.core import nodes from agent.core.state import AgentState from services.core.prompt_manager import PromptManager from services.core.sql_prompt_manager import SqlPromptManager from workflows.workflow_manager import WorkflowManager, WorkflowType class FakeResponse: def __init__(self, content: str, tool_calls=None): self.content = content self.tool_calls = tool_calls or [] class FakeModel: def invoke(self, messages): if len(messages) == 2 and getattr(messages[0], "content", "").startswith("You are a translation and normalization assistant"): return FakeResponse(messages[1].content) if len(messages) == 2 and "table prompt JSON" in getattr(messages[0], "content", ""): user_content = getattr(messages[1], "content", "") lowered = user_content.lower() if "query mode: topn" in lowered or "top 20" in lowered: return FakeResponse("SELECT ship_to_country as country, count(soid) as qty FROM dwd_ai.apbo_eta_ful WHERE data_flag = 'Newest' GROUP BY ship_to_country ORDER BY qty DESC LIMIT 20") if "query mode: aggregate" in lowered: return FakeResponse("SELECT region, count(soid) as qty FROM dwd_ai.apbo_eta_ful WHERE data_flag = 'Newest' GROUP BY region ORDER BY qty DESC") return FakeResponse("SELECT service_order_id FROM dwd_ai.apbo_eta_ful WHERE data_flag = 'Newest'") return FakeResponse("fallback") class FakeTemplateMatcher: def match(self, normalized_text: str): return { "table_name": "apbo_eta_ful", "candidates": [{"table_name": "apbo_eta_ful", "metadata": {"table": "apbo_eta_ful"}}], "raw": {"query": normalized_text}, } class FakeSqlPromptManager: def get_prompt(self, table_name: str): manager = SqlPromptManager() return manager.get_prompt(table_name) class EmptyTemplateMatcher: def match(self, normalized_text: str): return { "table_name": None, "candidates": [], "raw": {"query": normalized_text}, } class FirstTurnOnlyTemplateMatcher: def match(self, normalized_text: str): lowered = (normalized_text or "").lower() if "4020438779" in lowered or "query so" in lowered: return FakeTemplateMatcher().match(normalized_text) return EmptyTemplateMatcher().match(normalized_text) def test_prompt_manager_resolves_project_prompt_file(): manager = PromptManager() prompt_text = manager.get("system", "sql_mysql_select_only") assert "expert SQL generator" in prompt_text assert "Only output a single MySQL SELECT statement." in prompt_text def test_sql_prompt_manager_resolves_project_sql_prompt_dir(): manager = SqlPromptManager() prompt = manager.get_prompt("apbo_eta_ful") assert prompt is not None assert prompt["meta"]["data_source"] == "dwd_ai.apbo_eta_ful" def test_apbo_eta_ful_prompt_uses_identifier_specific_where_filters(): prompt = SqlPromptManager().get_prompt("apbo_eta_ful") rule_text = prompt["business_logic_rules"]["soid_or_service_order_id"] examples = prompt["examples"] assert "禁止使用WHERE soid" in rule_text assert "必须使用(service_order_id = 'xxx' or soid = 'xxx')" in rule_text assert "service_order_id" in examples["history_records_all_fields"]["sql"] assert " or soid " in examples["history_records_all_fields"]["sql"] assert " or soid " in examples["specific_fields_query"]["sql"] assert " or soid " in examples["newest_status_all_fields"]["sql"] def test_query_mode_and_plan_build_for_topn(): prompt = SqlPromptManager().get_prompt("apbo_eta_ful") state = AgentState( messages=[HumanMessage(content="CC为LT ,by country 查询 top 20")], context={"original_input": "CC为LT ,by country 查询 top 20", "normalized_input": "CC=LT by country top 20", "table_name": "apbo_eta_ful"}, ) state.sql_prompt = prompt state = nodes.classify_query_mode(state) state = nodes.build_sql_plan(state) assert state.query_mode == "topn" assert state.sql_plan["query_entities"]["top_n"] == 20 assert state.sql_plan["selected_table"] == "apbo_eta_ful" assert "top_n_rules" in state.sql_plan def test_query_mode_country_does_not_trigger_aggregate(): state = AgentState( messages=[HumanMessage(content="查询 AU country 的 DC premier stock backlog")], context={ "original_input": "查询 AU country 的 DC premier stock backlog", "normalized_input": "List of backlogs with values for AU Country DC premier stock", }, ) state = nodes.classify_query_mode(state) assert state.query_mode == "detail" def test_query_mode_explicit_aggregate_keyword_still_matches(): state = AgentState( messages=[HumanMessage(content="按 country 统计 backlog 数量")], context={ "original_input": "按 country 统计 backlog 数量", "normalized_input": "Count backlog quantity by country", }, ) state = nodes.classify_query_mode(state) assert state.query_mode == "aggregate" def test_query_mode_current_keyword_is_not_misclassified_as_topn(): state = AgentState( messages=[HumanMessage(content="汇总统计当前的bo数量")], context={ "original_input": "汇总统计当前的bo数量", "normalized_input": "Aggregate summary of the current backlog order count", }, ) state = nodes.classify_query_mode(state) assert state.query_mode == "aggregate" assert state.query_entities.get("top_n") is None def test_match_table_falls_back_to_config_default(monkeypatch): monkeypatch.setattr("agent.core.nodes.get_template_matcher", lambda: EmptyTemplateMatcher()) monkeypatch.setattr( "agent.core.nodes.Config.get_section", lambda section: {"default_table_name": "apbo_eta_ful"} if section == "ragflow" else {}, ) state = AgentState( messages=[HumanMessage(content="random question")], context={"original_input": "random question", "normalized_input": "random question"}, ) state = nodes.match_table(state) assert state.table_name == "apbo_eta_ful" assert state.context["table_match_fallback"] == "config_default" assert state.context["default_table_name"] == "apbo_eta_ful" def test_generate_response_falls_back_to_model_when_sr_api_result_empty(): class EmptyResultModel: def invoke(self, messages): return FakeResponse("未查询到符合条件的数据，请尝试调整筛选条件。") state = AgentState( messages=[HumanMessage(content="查询 AU 的 backlog")], context={ "original_input": "查询 AU 的 backlog", "normalized_input": "Query backlog for AU", "final_sql": "SELECT * FROM dwd_ai.apbo_eta_ful WHERE ship_to_country='AU'", "sr_api_result": '{"status_code": 200, "text": "{\\"code\\":\\"0\\",\\"data\\":[],\\"msg\\":\\"操作成功\\",\\"total\\":0}"}', }, ) state.final_sql = state.context["final_sql"] state.sr_api_result = state.context["sr_api_result"] state = nodes.generate_response(state, EmptyResultModel()) assert state.messages[-1].content.startswith("未查询到符合条件的数据") assert state.context["response_source"] == "model_empty_result_fallback" assert state.context["is_empty_result"] is True def test_conversation_agent_run_keeps_final_sql_in_context(monkeypatch): monkeypatch.setattr("agent.core.base_agent.create_chat_model", lambda model_section=None: FakeModel()) monkeypatch.setattr("agent.core.nodes.get_template_matcher", lambda: FakeTemplateMatcher()) monkeypatch.setattr("agent.core.nodes.get_sql_prompt_manager", lambda: FakeSqlPromptManager()) agent = ConversationAgent() result = agent.run("查询 SO 4020438779 的 eta 信息", skip_sr_api=True) assert result["context"]["table_name"] == "apbo_eta_ful" assert result["context"]["query_mode"] == "detail" assert result["context"]["final_sql"].startswith("SELECT") assert result["messages"][-1].content == result["context"]["final_sql"] def test_conversation_agent_requires_explicit_memory_for_follow_up(monkeypatch): monkeypatch.setattr("agent.core.base_agent.create_chat_model", lambda model_section=None: FakeModel()) monkeypatch.setattr("agent.core.nodes.get_template_matcher", lambda: FirstTurnOnlyTemplateMatcher()) monkeypatch.setattr("agent.core.nodes.get_sql_prompt_manager", lambda: FakeSqlPromptManager()) agent = ConversationAgent() first = agent.run("查询 SO 4020438779 的 eta 信息", skip_sr_api=True) second = agent.run("改成 by country top 20", skip_sr_api=True) third = agent.run( "改成 by country top 20", skip_sr_api=True, conversation_history=first["conversation_history"], last_context=first["context"].get("last_context"), ) assert first["context"]["table_name"] == "apbo_eta_ful" assert second["context"]["table_name"] == "apbo_eta_ful" assert second["context"]["table_match_fallback"] == "config_default" assert third["context"]["table_name"] == "apbo_eta_ful" assert third["context"]["table_match_fallback"] == "last_context" assert third["context"]["query_mode"] == "topn" assert "LIMIT 20" in third["context"]["final_sql"] def test_workflow_manager_follow_up_does_not_reuse_session_last_context(monkeypatch): monkeypatch.setattr("agent.core.base_agent.create_chat_model", lambda model_section=None: FakeModel()) monkeypatch.setattr("agent.core.nodes.get_template_matcher", lambda: FirstTurnOnlyTemplateMatcher()) monkeypatch.setattr("agent.core.nodes.get_sql_prompt_manager", lambda: FakeSqlPromptManager()) manager = WorkflowManager(enable_multi_turn=False) first = manager.execute_workflow( WorkflowType.CONVERSATION, "查询 SO 4020438779 的 eta 信息", session_id="cid-1", skip_sr_api=True, ) second = manager.execute_workflow( WorkflowType.CONVERSATION, "改成 by country top 20", session_id="cid-1", skip_sr_api=True, ) first_context = first["result"]["context"] second_context = second["result"]["context"] session_info = manager.get_session_info("cid-1") assert first_context["table_name"] == "apbo_eta_ful" assert second_context["table_name"] == "apbo_eta_ful" assert second_context["table_match_fallback"] == "config_default" assert second_context["query_mode"] == "topn" assert "LIMIT 20" in second_context["final_sql"] assert "last_context" not in session_info assert "conversation_history" not in session_info def test_workflow_manager_follow_up_reuses_session_last_context_when_enabled(monkeypatch): monkeypatch.setattr("agent.core.base_agent.create_chat_model", lambda model_section=None: FakeModel()) monkeypatch.setattr("agent.core.nodes.get_template_matcher", lambda: FirstTurnOnlyTemplateMatcher()) monkeypatch.setattr("agent.core.nodes.get_sql_prompt_manager", lambda: FakeSqlPromptManager()) manager = WorkflowManager(enable_multi_turn=True) first = manager.execute_workflow( WorkflowType.CONVERSATION, "查询 SO 4020438779 的 eta 信息", session_id="cid-enabled", skip_sr_api=True, ) second = manager.execute_workflow( WorkflowType.CONVERSATION, "改成 by country top 20", session_id="cid-enabled", skip_sr_api=True, ) first_context = first["result"]["context"] second_context = second["result"]["context"] session_info = manager.get_session_info("cid-enabled") assert first_context["table_name"] == "apbo_eta_ful" assert second_context["table_name"] == "apbo_eta_ful" assert second_context["table_match_fallback"] == "last_context" assert second_context["query_mode"] == "topn" assert "LIMIT 20" in second_context["final_sql"] assert session_info["last_context"]["table_name"] == "apbo_eta_ful" assert len(session_info["conversation_history"]) >= 4 def test_workflow_manager_isolates_conversation_memory_by_session(monkeypatch): monkeypatch.setattr("agent.core.base_agent.create_chat_model", lambda model_section=None: FakeModel()) monkeypatch.setattr("agent.core.nodes.get_template_matcher", lambda: FirstTurnOnlyTemplateMatcher()) monkeypatch.setattr("agent.core.nodes.get_sql_prompt_manager", lambda: FakeSqlPromptManager()) manager = WorkflowManager(enable_multi_turn=False) manager.execute_workflow( WorkflowType.CONVERSATION, "查询 SO 4020438779 的 eta 信息", session_id="session-a", skip_sr_api=True, ) second = manager.execute_workflow( WorkflowType.CONVERSATION, "改成 by country top 20", session_id="session-b", skip_sr_api=True, ) assert second["result"]["context"]["table_name"] == "apbo_eta_ful" assert second["result"]["context"]["table_match_fallback"] == "config_default" def test_conversation_agent_run_is_silent_without_debug_prints(monkeypatch, capsys): monkeypatch.setattr("agent.core.base_agent.create_chat_model", lambda model_section=None: FakeModel()) monkeypatch.setattr("agent.core.nodes.get_template_matcher", lambda: FakeTemplateMatcher()) monkeypatch.setattr("agent.core.nodes.get_sql_prompt_manager", lambda: FakeSqlPromptManager()) agent = ConversationAgent() agent.run("查询 SO 4020438779 的 eta 信息", skip_sr_api=True) captured = capsys.readouterr() assert captured.out == "" def test_conversation_agent_run_prints_node_trace_when_enabled(monkeypatch, capsys): monkeypatch.setattr("agent.core.base_agent.create_chat_model", lambda model_section=None: FakeModel()) monkeypatch.setattr("agent.core.nodes.get_template_matcher", lambda: FakeTemplateMatcher()) monkeypatch.setattr("agent.core.nodes.get_sql_prompt_manager", lambda: FakeSqlPromptManager()) agent = ConversationAgent() agent.run("查询 SO 4020438779 的 eta 信息", skip_sr_api=True, debug_node_trace=True) captured = capsys.readouterr() assert "[process_input][in]" in captured.out assert "[generate_response][out] source=final_sql" in captured.out def test_generate_sql_trace_outputs_full_sql_when_enabled(capsys): class LongSqlModel: def invoke(self, messages): return FakeResponse("SELECT " + ", ".join(f"col_{idx}" for idx in range(150)) + " FROM dwd_ai.apbo_eta_ful") state = AgentState( messages=[HumanMessage(content="查询 long sql")], context={ "original_input": "查询 long sql", "normalized_input": "Query long sql", "table_name": "apbo_eta_ful", "sql_plan": {"selected_table": "apbo_eta_ful"}, "debug_node_trace": True, }, ) state.sql_prompt = SqlPromptManager().get_prompt("apbo_eta_ful") state = nodes.generate_sql(state, LongSqlModel()) captured = capsys.readouterr() assert "[generate_sql][out] sql=" in captured.out assert "col_149" in captured.out assert "..." not in captured.out.split("[generate_sql][out] sql=", 1)[1]