diff --git a/api/core/agent/base_agent_runner.py b/api/core/agent/base_agent_runner.py index 806f7c6590f..c8bf7173161 100644 --- a/api/core/agent/base_agent_runner.py +++ b/api/core/agent/base_agent_runner.py @@ -1,8 +1,9 @@ import json import logging import uuid +from collections import Counter from decimal import Decimal -from typing import Union, cast +from typing import Any, Union, cast from sqlalchemy import func, select from sqlalchemy.orm import Session @@ -48,6 +49,34 @@ logger = logging.getLogger(__name__) _file_access_controller = DatabaseFileAccessController() +def _select_tool_occurrence(value: Any, occurrence: int, occurrences: int) -> Any: + """ + Pick one call's value out of a persisted agent thought payload. + + A tool called several times in one turn stores one value per call, in call + order. Records written before those calls were kept apart store a single + value for the tool name, and every occurrence replays it — the behaviour + those records were written with. + + `observation` values are always `str`: `ToolEngine.agent_invoke` is typed + `-> tuple[str, list[str], ToolInvokeMeta]` and the runners store element 0, + so a list there is not a shape any writer produces and the length check is + defensive. `tool_input` values are `json.loads` of the model's `arguments` + with no shape check, so a legacy list-valued input is possible in principle; + on a length collision the rule reads per call, not whole — + `test_a_legacy_list_of_matching_length_is_read_per_call_not_whole` pins it. + + The same function is defined, identically, in `models/model.py`. + The duplication is deliberate: `models/` importing from `core/agent/` is the + worse layering trade and the reverse is odd, so neither copy is in the wrong + place and neither should move. A change to this rule must be applied in both. + """ + if occurrences > 1 and isinstance(value, list) and len(value) == occurrences: + return value[occurrence] + + return value + + class BaseAgentRunner(AppRunner): def __init__( self, @@ -408,7 +437,11 @@ class BaseAgentRunner(AppRunner): else: tool_responses = dict.fromkeys(tool_names, observation_payload) + tool_occurrences = Counter(tool_names) + seen_tools: Counter[str] = Counter() for tool in tool_names: + occurrence = seen_tools[tool] + seen_tools[tool] += 1 # generate a uuid for tool call tool_call_id = str(uuid.uuid4()) tool_calls.append( @@ -417,13 +450,21 @@ class BaseAgentRunner(AppRunner): type="function", function=AssistantPromptMessage.ToolCall.ToolCallFunction( name=tool, - arguments=json.dumps(tool_inputs.get(tool, {})), + arguments=json.dumps( + _select_tool_occurrence( + tool_inputs.get(tool, {}), occurrence, tool_occurrences[tool] + ) + ), ), ) ) tool_call_response.append( ToolPromptMessage( - content=tool_responses.get(tool, agent_thought.observation), + content=_select_tool_occurrence( + tool_responses.get(tool, agent_thought.observation), + occurrence, + tool_occurrences[tool], + ), name=tool, tool_call_id=tool_call_id, ) diff --git a/api/core/agent/fc_agent_runner.py b/api/core/agent/fc_agent_runner.py index 78980f0d943..9c83d660eb0 100644 --- a/api/core/agent/fc_agent_runner.py +++ b/api/core/agent/fc_agent_runner.py @@ -43,6 +43,21 @@ _FILE_PREVIEW_ID_PATTERN = re.compile(r"/files/([a-fA-F0-9-]{36})/file-preview") _KNOWLEDGE_RETRIEVAL_PROMPT_NAME = "knowledge_retrieval" +def _group_by_tool_name(items: list[tuple[str, Any]]) -> dict[str, Any]: + """ + Group per-call values by tool name for the agent thought record. + + A tool called once in a turn keeps its bare value, which is the shape every + record written so far uses. A tool called more than once keeps every value + in call order, so a repeated tool no longer discards all but the last one. + """ + grouped: dict[str, list[Any]] = {} + for tool_name, value in items: + grouped.setdefault(tool_name, []).append(value) + + return {tool_name: values[0] if len(values) == 1 else values for tool_name, values in grouped.items()} + + class FunctionCallAgentRunner(BaseAgentRunner): def _build_dataset_tool_image_contents( self, session: Session, tool_response: str, tool_instance: Any @@ -201,13 +216,12 @@ class FunctionCallAgentRunner(BaseAgentRunner): function_call_state = True tool_calls.extend(self.extract_tool_calls(chunk) or []) tool_call_names = ";".join([tool_call[1] for tool_call in tool_calls]) + grouped_inputs = _group_by_tool_name([(tool_call[1], tool_call[2]) for tool_call in tool_calls]) try: - tool_call_inputs = json.dumps( - {tool_call[1]: tool_call[2] for tool_call in tool_calls}, ensure_ascii=False - ) + tool_call_inputs = json.dumps(grouped_inputs, ensure_ascii=False) except TypeError: # fallback: force ASCII to handle non-serializable objects - tool_call_inputs = json.dumps({tool_call[1]: tool_call[2] for tool_call in tool_calls}) + tool_call_inputs = json.dumps(grouped_inputs) if chunk.delta.message and chunk.delta.message.content: if isinstance(chunk.delta.message.content, list): @@ -228,13 +242,12 @@ class FunctionCallAgentRunner(BaseAgentRunner): function_call_state = True tool_calls.extend(self.extract_blocking_tool_calls(result) or []) tool_call_names = ";".join([tool_call[1] for tool_call in tool_calls]) + grouped_inputs = _group_by_tool_name([(tool_call[1], tool_call[2]) for tool_call in tool_calls]) try: - tool_call_inputs = json.dumps( - {tool_call[1]: tool_call[2] for tool_call in tool_calls}, ensure_ascii=False - ) + tool_call_inputs = json.dumps(grouped_inputs, ensure_ascii=False) except TypeError: # fallback: force ASCII to handle non-serializable objects - tool_call_inputs = json.dumps({tool_call[1]: tool_call[2] for tool_call in tool_calls}) + tool_call_inputs = json.dumps(grouped_inputs) if result.usage: increase_usage(llm_usage, result.usage) @@ -380,13 +393,18 @@ class FunctionCallAgentRunner(BaseAgentRunner): tool_name="", tool_input="", thought="", - tool_invoke_meta={ - tool_response["tool_call_name"]: tool_response["meta"] for tool_response in tool_responses - }, - observation={ - tool_response["tool_call_name"]: tool_response["tool_response"] - for tool_response in tool_responses - }, + tool_invoke_meta=_group_by_tool_name( + [ + (str(tool_response["tool_call_name"]), tool_response["meta"]) + for tool_response in tool_responses + ] + ), + observation=_group_by_tool_name( + [ + (str(tool_response["tool_call_name"]), tool_response["tool_response"]) + for tool_response in tool_responses + ] + ), answer="", messages_ids=message_file_ids, ) diff --git a/api/core/ops/ops_trace_manager.py b/api/core/ops/ops_trace_manager.py index 0066d035e54..e49bd518ec0 100644 --- a/api/core/ops/ops_trace_manager.py +++ b/api/core/ops/ops_trace_manager.py @@ -1197,6 +1197,12 @@ class TraceTask: if tool_name in agent_thought.tools: created_time = agent_thought.created_at tool_meta_data = agent_thought.tool_meta.get(tool_name, {}) + if isinstance(tool_meta_data, list): + # a tool called several times in one turn keeps one meta per + # call; this trace covers the tool, not one call of it, so it + # takes the last — the only one a record written before those + # calls were kept apart had + tool_meta_data = tool_meta_data[-1] if tool_meta_data else {} tool_config = tool_meta_data.get("tool_config", {}) time_cost = tool_meta_data.get("time_cost", 0) end_time = created_time + timedelta(seconds=time_cost) diff --git a/api/models/model.py b/api/models/model.py index 0cad0bb3feb..bdbef220400 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -3,6 +3,7 @@ from __future__ import annotations import json import re import uuid +from collections import Counter from collections.abc import Callable, Mapping, Sequence from datetime import datetime from decimal import Decimal @@ -2523,6 +2524,34 @@ class MessageChain(TypeBase): ) +def _select_tool_occurrence(value: Any, occurrence: int, occurrences: int) -> Any: + """ + Pick one call's value out of a persisted agent thought payload. + + A tool called several times in one turn stores one value per call, in call + order. Records written before those calls were kept apart store a single + value for the tool name, and every occurrence replays it — the behaviour + those records were written with. + + `observation` values are always `str`: `ToolEngine.agent_invoke` is typed + `-> tuple[str, list[str], ToolInvokeMeta]` and the runners store element 0, + so a list there is not a shape any writer produces and the length check is + defensive. `tool_input` values are `json.loads` of the model's `arguments` + with no shape check, so a legacy list-valued input is possible in principle; + on a length collision the rule reads per call, not whole — + `test_a_legacy_list_of_matching_length_is_read_per_call_not_whole` pins it. + + The same function is defined, identically, in `core/agent/base_agent_runner.py`. + The duplication is deliberate: `models/` importing from `core/agent/` is the + worse layering trade and the reverse is odd, so neither copy is in the wrong + place and neither should move. A change to this rule must be applied in both. + """ + if occurrences > 1 and isinstance(value, list) and len(value) == occurrences: + return value[occurrence] + + return value + + class MessageAgentThought(TypeBase): __tablename__ = "message_agent_thoughts" __table_args__ = ( @@ -2644,6 +2673,38 @@ class MessageAgentThought(TypeBase): else: return {} + def _per_call(self, values_by_tool: dict[str, Any]) -> list[Any]: + """ + Spread a tool-name-keyed payload over the turn's calls, in call order. + + One entry per entry in ``tools``, so a tool called twice gets two, each + holding that call's own value. A name whose stored value is not one + value per call — every record written before those calls were kept + apart — gives the same value to each of its calls, which is what those + records mean. + """ + occurrences = Counter(self.tools) + seen: Counter[str] = Counter() + per_call: list[Any] = [] + for tool in self.tools: + occurrence = seen[tool] + seen[tool] += 1 + per_call.append(_select_tool_occurrence(values_by_tool.get(tool, {}), occurrence, occurrences[tool])) + + return per_call + + @property + def tool_inputs_per_call(self) -> list[Any]: + return self._per_call(self.tool_inputs_dict) + + @property + def tool_outputs_per_call(self) -> list[Any]: + return self._per_call(self.tool_outputs_dict) + + @property + def tool_metas_per_call(self) -> list[Any]: + return self._per_call(self.tool_meta) + class DatasetRetrieverResource(TypeBase): __tablename__ = "dataset_retriever_resources" diff --git a/api/services/agent_service.py b/api/services/agent_service.py index e065e1f978b..8671aabc11e 100644 --- a/api/services/agent_service.py +++ b/api/services/agent_service.py @@ -95,16 +95,16 @@ class AgentService: for agent_thought in agent_thoughts: tools = agent_thought.tools tool_labels = agent_thought.tool_labels - tool_meta = agent_thought.tool_meta - tool_inputs = agent_thought.tool_inputs_dict - tool_outputs = agent_thought.tool_outputs_dict or {} + # one entry per call, in call order, so a tool called twice in a turn + # gets a log entry per call instead of one call's data twice + tool_inputs = agent_thought.tool_inputs_per_call + tool_outputs = agent_thought.tool_outputs_per_call + tool_metas = agent_thought.tool_metas_per_call tool_calls = [] - for tool in tools: - tool_name = tool + for tool_name, tool_input, tool_output, tool_meta_data in zip( + tools, tool_inputs, tool_outputs, tool_metas, strict=True + ): tool_label = tool_labels.get(tool_name, tool_name) - tool_input = tool_inputs.get(tool_name, {}) - tool_output = tool_outputs.get(tool_name, {}) - tool_meta_data = tool_meta.get(tool_name, {}) tool_config = tool_meta_data.get("tool_config", {}) tool_provider_type = tool_config.get("tool_provider_type", "") tool_provider_id = tool_config.get("tool_provider", "") diff --git a/api/tests/unit_tests/core/agent/test_base_agent_runner.py b/api/tests/unit_tests/core/agent/test_base_agent_runner.py index 7411164512c..b7fd1037fdc 100644 --- a/api/tests/unit_tests/core/agent/test_base_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_base_agent_runner.py @@ -668,3 +668,116 @@ class TestBaseAgentRunnerCoverage: result = runner.organize_agent_history([], session=mock_db_session) assert any(isinstance(item, module.ToolPromptMessage) for item in result) + + +# ========================================================== +# organize_agent_history — a tool called more than once in a turn +# ========================================================== + + +class TestOrganizeHistoryRepeatedTools: + def _replay(self, runner: BaseAgentRunner, mock_db_session, mocker: MockerFixture, thought): + msg = mocker.MagicMock(id="m_repeat", agent_thoughts=[thought], answer=None, app_model_config=None) + msg.agent_thoughts_with_session.return_value = [thought] + + mock_db_session.execute.return_value.scalars.return_value.all.return_value = [msg] + mocker.patch.object(module, "extract_thread_messages", return_value=[msg]) + mocker.patch.object( + runner, + "organize_agent_user_prompt", + return_value=module.UserPromptMessage(content="user"), + ) + + result = runner.organize_agent_history([], session=mock_db_session) + + assistant = next(item for item in result if isinstance(item, module.AssistantPromptMessage) and item.tool_calls) + responses = [item for item in result if isinstance(item, module.ToolPromptMessage)] + return assistant, responses + + def test_repeated_tool_replays_each_call_separately( + self, runner: BaseAgentRunner, mock_db_session, mocker: MockerFixture + ): + thought = mocker.MagicMock( + tool="search;search", + tool_input=json.dumps({"search": [{"q": "first"}, {"q": "second"}]}), + observation=json.dumps({"search": ["first result", "second result"]}), + thought="thinking", + ) + + assistant, responses = self._replay(runner, mock_db_session, mocker, thought) + + assert [call.function.name for call in assistant.tool_calls] == ["search", "search"] + assert [json.loads(call.function.arguments) for call in assistant.tool_calls] == [ + {"q": "first"}, + {"q": "second"}, + ] + assert [response.content for response in responses] == ["first result", "second result"] + # each replayed response is tied to its own call + assert [call.id for call in assistant.tool_calls] == [response.tool_call_id for response in responses] + assert len({call.id for call in assistant.tool_calls}) == 2 + + def test_legacy_record_replays_unchanged(self, runner: BaseAgentRunner, mock_db_session, mocker: MockerFixture): + # exactly the shape written before repeated calls were kept apart: one + # value per tool name, the last call having overwritten the first + thought = mocker.MagicMock( + tool="search;search", + tool_input=json.dumps({"search": {"q": "second"}}), + observation=json.dumps({"search": "second result"}), + thought="thinking", + ) + + assistant, responses = self._replay(runner, mock_db_session, mocker, thought) + + assert [json.loads(call.function.arguments) for call in assistant.tool_calls] == [ + {"q": "second"}, + {"q": "second"}, + ] + assert [response.content for response in responses] == ["second result", "second result"] + + def test_single_call_replays_unchanged(self, runner: BaseAgentRunner, mock_db_session, mocker: MockerFixture): + thought = mocker.MagicMock( + tool="search", + tool_input=json.dumps({"search": {"q": "only"}}), + observation=json.dumps({"search": "only result"}), + thought="thinking", + ) + + assistant, responses = self._replay(runner, mock_db_session, mocker, thought) + + assert len(assistant.tool_calls) == 1 + assert json.loads(assistant.tool_calls[0].function.arguments) == {"q": "only"} + assert [response.content for response in responses] == ["only result"] + + def test_distinct_tools_replay_unchanged(self, runner: BaseAgentRunner, mock_db_session, mocker: MockerFixture): + thought = mocker.MagicMock( + tool="search;calculator", + tool_input=json.dumps({"search": {"q": "a"}, "calculator": {"expr": "1+1"}}), + observation=json.dumps({"search": "search result", "calculator": "2"}), + thought="thinking", + ) + + assistant, responses = self._replay(runner, mock_db_session, mocker, thought) + + assert [call.function.name for call in assistant.tool_calls] == ["search", "calculator"] + assert [response.content for response in responses] == ["search result", "2"] + + def test_a_legacy_list_of_matching_length_is_read_per_call_not_whole( + self, runner: BaseAgentRunner, mock_db_session, mocker: MockerFixture + ): + # the other half of the same decision. Length is the only signal the + # reader has, so a single stored value that is itself a list as long as + # the call count is indistinguishable from one value per call, and is + # read as one value per call. A legacy row whose one value happened to + # be a two-element list is therefore split across the two calls instead + # of replayed whole. + thought = mocker.MagicMock( + tool="search;search", + tool_input=json.dumps({"search": ["a", "b"]}), + observation=json.dumps({"search": ["x", "y"]}), + thought="thinking", + ) + + assistant, responses = self._replay(runner, mock_db_session, mocker, thought) + + assert [json.loads(call.function.arguments) for call in assistant.tool_calls] == ["a", "b"] + assert [response.content for response in responses] == ["x", "y"] diff --git a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py index ea8c3b87c42..33ff0c237f1 100644 --- a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py @@ -574,3 +574,133 @@ class TestRunMethod: with pytest.raises(AgentMaxIterationError): list(runner.run(runner.session, message, "query")) + + +# ============================== +# Repeated Tool Calls In One Turn +# ============================== + + +def _make_tool_call(call_id: str, name: str, arguments: dict[str, Any]) -> MagicMock: + tool_call = MagicMock() + tool_call.id = call_id + tool_call.function.name = name + tool_call.function.arguments = json.dumps(arguments) + return tool_call + + +def _run_one_turn(runner: FunctionCallAgentRunner, mocker: MockerFixture, tool_calls, responses, metas=None): + """Drive one tool-calling iteration followed by a plain answer.""" + result = DummyResult(message=DummyMessage(content="", tool_calls=tool_calls), usage=build_usage()) + final_result = DummyResult(message=DummyMessage(content="done", tool_calls=[]), usage=build_usage()) + runner.model_instance.invoke_llm.side_effect = [result, final_result] + + names = {tool_call.function.name for tool_call in tool_calls} + prompt_tools = [] + instances = {} + for name in names: + prompt_tool = MagicMock() + prompt_tool.name = name + prompt_tools.append(prompt_tool) + instances[name] = MagicMock() + runner._init_prompt_tools.return_value = (instances, prompt_tools) + + invoke_metas = [] + for meta in metas or [{"ok": True} for _ in responses]: + tool_invoke_meta = MagicMock() + tool_invoke_meta.to_dict.return_value = meta + invoke_metas.append(tool_invoke_meta) + mocker.patch( + "core.agent.fc_agent_runner.ToolEngine.agent_invoke", + side_effect=[ + (response, [], tool_invoke_meta) for response, tool_invoke_meta in zip(responses, invoke_metas, strict=True) + ], + ) + + list(runner.run(runner.session, MagicMock(id="m1"), "query")) + + +def _saved_tool_input(runner: FunctionCallAgentRunner) -> dict[str, Any]: + for call in runner.save_agent_thought.call_args_list: + if call.kwargs.get("tool_name"): + return json.loads(call.kwargs["tool_input"]) + raise AssertionError("no agent thought was saved with tool input") + + +def _saved_observation(runner: FunctionCallAgentRunner) -> dict[str, Any]: + for call in runner.save_agent_thought.call_args_list: + observation = call.kwargs.get("observation") + if observation: + return observation + raise AssertionError("no agent thought was saved with an observation") + + +def _saved_tool_invoke_meta(runner: FunctionCallAgentRunner) -> dict[str, Any]: + for call in runner.save_agent_thought.call_args_list: + tool_invoke_meta = call.kwargs.get("tool_invoke_meta") + if tool_invoke_meta: + return tool_invoke_meta + raise AssertionError("no agent thought was saved with tool invoke meta") + + +class TestRepeatedToolCalls: + def test_repeated_tool_keeps_every_input_and_observation( + self, runner: FunctionCallAgentRunner, mocker: MockerFixture + ): + _run_one_turn( + runner, + mocker, + tool_calls=[ + _make_tool_call("1", "search", {"q": "first"}), + _make_tool_call("2", "search", {"q": "second"}), + ], + responses=["first result", "second result"], + ) + + assert _saved_tool_input(runner) == {"search": [{"q": "first"}, {"q": "second"}]} + assert _saved_observation(runner) == {"search": ["first result", "second result"]} + + def test_repeated_tool_keeps_every_invoke_meta(self, runner: FunctionCallAgentRunner, mocker: MockerFixture): + # the agent log reads a provider, a duration and an error out of the meta + # for each call it displays, so a repeated tool needs one meta per call + _run_one_turn( + runner, + mocker, + tool_calls=[ + _make_tool_call("1", "search", {"q": "first"}), + _make_tool_call("2", "search", {"q": "second"}), + ], + responses=["first result", "second result"], + metas=[{"time_cost": 1}, {"time_cost": 2}], + ) + + assert _saved_tool_invoke_meta(runner) == {"search": [{"time_cost": 1}, {"time_cost": 2}]} + + def test_single_call_keeps_the_legacy_shape(self, runner: FunctionCallAgentRunner, mocker: MockerFixture): + _run_one_turn( + runner, + mocker, + tool_calls=[_make_tool_call("1", "search", {"q": "only"})], + responses=["only result"], + metas=[{"time_cost": 1}], + ) + + assert _saved_tool_input(runner) == {"search": {"q": "only"}} + assert _saved_observation(runner) == {"search": "only result"} + assert _saved_tool_invoke_meta(runner) == {"search": {"time_cost": 1}} + + def test_distinct_tools_keep_the_legacy_shape(self, runner: FunctionCallAgentRunner, mocker: MockerFixture): + _run_one_turn( + runner, + mocker, + tool_calls=[ + _make_tool_call("1", "search", {"q": "a"}), + _make_tool_call("2", "calculator", {"expr": "1+1"}), + ], + responses=["search result", "2"], + metas=[{"time_cost": 1}, {"time_cost": 2}], + ) + + assert _saved_tool_input(runner) == {"search": {"q": "a"}, "calculator": {"expr": "1+1"}} + assert _saved_observation(runner) == {"search": "search result", "calculator": "2"} + assert _saved_tool_invoke_meta(runner) == {"search": {"time_cost": 1}, "calculator": {"time_cost": 2}} diff --git a/api/tests/unit_tests/core/ops/test_ops_trace_manager.py b/api/tests/unit_tests/core/ops/test_ops_trace_manager.py index 3c3011597f0..77d65062f50 100644 --- a/api/tests/unit_tests/core/ops/test_ops_trace_manager.py +++ b/api/tests/unit_tests/core/ops/test_ops_trace_manager.py @@ -538,6 +538,28 @@ def test_tool_trace_reads_real_message_file(monkeypatch: pytest.MonkeyPatch, dat assert result.message_file_data.id == file.id +def test_tool_trace_reads_the_last_call_of_a_repeated_tool(monkeypatch: pytest.MonkeyPatch, database: Session) -> None: + # a tool called twice in one turn keeps one meta per call; the trace covers + # the tool rather than one call of it, so it reads the last — the only one a + # record written before those calls were kept apart held + thought = SimpleNamespace( + tools=["tool-a", "tool-a"], + created_at=datetime(2025, 2, 20, 12, 1), + tool_meta={ + "tool-a": [ + {"tool_config": {}, "time_cost": 3, "error": "", "tool_parameters": {"q": "first"}}, + {"tool_config": {}, "time_cost": 5, "error": "", "tool_parameters": {"q": "second"}}, + ] + }, + ) + monkeypatch.setattr(module, "get_message_data", lambda _message_id: _message_data(agent_thoughts=[thought])) + result = TraceTask(trace_type=TraceTaskName.TOOL_TRACE).tool_trace( + "message-1", {"start": 1, "end": 2}, tool_name="tool-a", tool_inputs={}, tool_outputs="result" + ) + assert result.time_cost == 5 + assert result.tool_parameters == {"q": "second"} + + def test_node_execution_trace_resolves_real_message_by_conversation_and_run( trace_environment: None, database: Session ) -> None: diff --git a/api/tests/unit_tests/models/test_model.py b/api/tests/unit_tests/models/test_model.py index bb3206713d7..8810a910b9c 100644 --- a/api/tests/unit_tests/models/test_model.py +++ b/api/tests/unit_tests/models/test_model.py @@ -1,4 +1,5 @@ import importlib +import json import types from unittest.mock import MagicMock, patch @@ -6,7 +7,8 @@ import pytest from core.workflow.file_reference import build_file_reference from graphon.file import FILE_MODEL_IDENTITY, FileTransferMethod -from models.model import Conversation, Message +from models.enums import CreatorUserRole +from models.model import Conversation, Message, MessageAgentThought @pytest.fixture(autouse=True) @@ -138,3 +140,134 @@ def test_message_inputs_resolve_file_tenant_with_caller_session() -> None: assert inputs["file"] == "tenant-1" session.scalar.assert_called_once() + + +# ========================================================== +# MessageAgentThought — one payload per call, in call order +# ========================================================== + + +def _agent_thought(*, tool: str, tool_input: str, observation: str, tool_meta_str: str) -> MessageAgentThought: + return MessageAgentThought( + message_id="message-1", + position=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + tool=tool, + tool_input=tool_input, + observation=observation, + tool_meta_str=tool_meta_str, + ) + + +# the four shapes an agent log can be built from: a record written before +# repeated calls were kept apart, and one written after, each with the tool +# called once and called twice. The legacy pair are written as literals — they +# are rows that exist and cannot be migrated. +LEGACY_SINGLE = _agent_thought( + tool="search", + tool_input=json.dumps({"search": {"q": "only"}}), + observation=json.dumps({"search": "only result"}), + tool_meta_str=json.dumps({"search": {"time_cost": 1}}), +) +LEGACY_REPEATED = _agent_thought( + tool="search;search", + tool_input=json.dumps({"search": {"q": "second"}}), + observation=json.dumps({"search": "second result"}), + tool_meta_str=json.dumps({"search": {"time_cost": 2}}), +) +NEW_SINGLE = _agent_thought( + tool="search", + tool_input=json.dumps({"search": {"q": "only"}}), + observation=json.dumps({"search": "only result"}), + tool_meta_str=json.dumps({"search": {"time_cost": 1}}), +) +NEW_REPEATED = _agent_thought( + tool="search;search", + tool_input=json.dumps({"search": [{"q": "first"}, {"q": "second"}]}), + observation=json.dumps({"search": ["first result", "second result"]}), + tool_meta_str=json.dumps({"search": [{"time_cost": 1}, {"time_cost": 2}]}), +) + + +def test_legacy_single_call_reads_one_value_per_payload(): + assert LEGACY_SINGLE.tool_inputs_per_call == [{"q": "only"}] + assert LEGACY_SINGLE.tool_outputs_per_call == ["only result"] + assert LEGACY_SINGLE.tool_metas_per_call == [{"time_cost": 1}] + + +def test_legacy_repeated_call_replays_the_surviving_value_for_each_call(): + # the row only ever held one call's data; both entries show it, exactly as + # they do without this change + assert LEGACY_REPEATED.tool_inputs_per_call == [{"q": "second"}, {"q": "second"}] + assert LEGACY_REPEATED.tool_outputs_per_call == ["second result", "second result"] + assert LEGACY_REPEATED.tool_metas_per_call == [{"time_cost": 2}, {"time_cost": 2}] + + +def test_new_single_call_reads_identically_to_a_legacy_single_call(): + assert NEW_SINGLE.tool_inputs_per_call == LEGACY_SINGLE.tool_inputs_per_call + assert NEW_SINGLE.tool_outputs_per_call == LEGACY_SINGLE.tool_outputs_per_call + assert NEW_SINGLE.tool_metas_per_call == LEGACY_SINGLE.tool_metas_per_call + + +def test_new_repeated_call_reads_each_call_separately(): + assert NEW_REPEATED.tool_inputs_per_call == [{"q": "first"}, {"q": "second"}] + assert NEW_REPEATED.tool_outputs_per_call == ["first result", "second result"] + assert NEW_REPEATED.tool_metas_per_call == [{"time_cost": 1}, {"time_cost": 2}] + + +def test_distinct_tools_read_one_value_each(): + thought = _agent_thought( + tool="search;calculator", + tool_input=json.dumps({"search": {"q": "a"}, "calculator": {"expr": "1+1"}}), + observation=json.dumps({"search": "search result", "calculator": "2"}), + tool_meta_str=json.dumps({"search": {"time_cost": 1}, "calculator": {"time_cost": 2}}), + ) + + assert thought.tool_inputs_per_call == [{"q": "a"}, {"expr": "1+1"}] + assert thought.tool_outputs_per_call == ["search result", "2"] + assert thought.tool_metas_per_call == [{"time_cost": 1}, {"time_cost": 2}] + + +def test_a_stored_list_that_is_not_one_value_per_call_is_replayed_whole(): + # the reader tells a per-call list from a single call's list value by length + # alone; a length that does not match the call count is not per-call data + thought = _agent_thought( + tool="search;search", + tool_input=json.dumps({"search": ["a", "b", "c"]}), + observation=json.dumps({"search": ["x", "y", "z"]}), + tool_meta_str=json.dumps({"search": {"time_cost": 1}}), + ) + + assert thought.tool_inputs_per_call == [["a", "b", "c"], ["a", "b", "c"]] + assert thought.tool_outputs_per_call == [["x", "y", "z"], ["x", "y", "z"]] + + +def test_a_legacy_list_of_matching_length_is_read_per_call_not_whole(): + # the other half of the same decision. Length is the only signal the reader + # has, so a single stored value that is itself a list as long as the call + # count is indistinguishable from one value per call, and is read as one + # value per call. A legacy row whose one value happened to be a two-element + # list is therefore split across the two calls instead of replayed whole. + thought = _agent_thought( + tool="search;search", + tool_input=json.dumps({"search": ["a", "b"]}), + observation=json.dumps({"search": ["x", "y"]}), + tool_meta_str=json.dumps({"search": {"time_cost": 1}}), + ) + + assert thought.tool_inputs_per_call == ["a", "b"] + assert thought.tool_outputs_per_call == ["x", "y"] + + +def test_a_tool_missing_from_the_payload_reads_empty(): + thought = _agent_thought( + tool="search;calculator", + tool_input=json.dumps({"search": {"q": "a"}}), + observation=json.dumps({"search": "search result"}), + tool_meta_str="{}", + ) + + assert thought.tool_inputs_per_call == [{"q": "a"}, {}] + assert thought.tool_outputs_per_call == ["search result", {}] + assert thought.tool_metas_per_call == [{}, {}] diff --git a/api/tests/unit_tests/services/test_agent_service.py b/api/tests/unit_tests/services/test_agent_service.py new file mode 100644 index 00000000000..e4e7b04c0ee --- /dev/null +++ b/api/tests/unit_tests/services/test_agent_service.py @@ -0,0 +1,128 @@ +import json +from datetime import UTC, datetime +from typing import TypedDict, cast +from unittest.mock import MagicMock, patch + +import pytest + +from models.account import Account +from models.enums import CreatorUserRole +from models.model import MessageAgentThought +from services.agent_service import AgentService + + +class _Iteration(TypedDict): + tool_calls: list[dict[str, object]] + + +class _AgentLogs(TypedDict): + iterations: list[_Iteration] + + +def _agent_thought(*, tool: str, tool_input: str, observation: str, tool_meta_str: str) -> MessageAgentThought: + thought = MessageAgentThought( + message_id="message-1", + position=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + tool=tool, + tool_input=tool_input, + observation=observation, + tool_meta_str=tool_meta_str, + tool_labels_str=json.dumps({"search": {"en_US": "Search"}}), + ) + thought.created_at = datetime(2026, 8, 5, tzinfo=UTC) + thought.tokens = 10 + return thought + + +def _get_agent_logs(agent_thought: MessageAgentThought) -> _AgentLogs: + app_model = MagicMock(id="app-1", tenant_id="tenant-1") + session = MagicMock() + conversation = MagicMock(from_end_user_id=None, from_account_id="account-1") + message = MagicMock(provider_response_latency=1.5, answer_tokens=1, message_tokens=2) + message.agent_thoughts_with_session.return_value = [agent_thought] + session.scalar.side_effect = [conversation, message, "Executor"] + + with ( + patch("services.agent_service.current_user", MagicMock(spec=Account, timezone="UTC")), + patch("services.agent_service.load_annotation_reply_config", return_value=None), + patch("services.agent_service.AgentConfigManager.convert", return_value=MagicMock(tools=[])), + patch("services.agent_service.ToolManager.get_tool_icon", return_value="icon"), + ): + return cast(_AgentLogs, AgentService.get_agent_logs(app_model, "conversation-1", "message-1", session)) + + +@pytest.fixture +def legacy_repeated() -> MessageAgentThought: + # written before repeated calls were kept apart: the second call's data + # overwrote the first, and the row holds one value for the tool name + return _agent_thought( + tool="search;search", + tool_input=json.dumps({"search": {"q": "second"}}), + observation=json.dumps({"search": "second result"}), + tool_meta_str=json.dumps({"search": {"time_cost": 2, "tool_config": {"tool_provider_type": "builtin"}}}), + ) + + +def test_repeated_tool_renders_one_log_entry_per_call() -> None: + thought = _agent_thought( + tool="search;search", + tool_input=json.dumps({"search": [{"q": "first"}, {"q": "second"}]}), + observation=json.dumps({"search": ["first result", "second result"]}), + tool_meta_str=json.dumps( + { + "search": [ + {"time_cost": 1, "tool_config": {"tool_provider_type": "builtin", "tool_provider": "a"}}, + {"time_cost": 2, "tool_config": {"tool_provider_type": "api", "tool_provider": "b"}}, + ] + } + ), + ) + + tool_calls = _get_agent_logs(thought)["iterations"][0]["tool_calls"] + + assert [call["tool_name"] for call in tool_calls] == ["search", "search"] + assert [call["tool_input"] for call in tool_calls] == [{"q": "first"}, {"q": "second"}] + assert [call["tool_output"] for call in tool_calls] == ["first result", "second result"] + assert [call["time_cost"] for call in tool_calls] == [1, 2] + + +def test_legacy_repeated_tool_renders_exactly_as_it_did(legacy_repeated: MessageAgentThought) -> None: + tool_calls = _get_agent_logs(legacy_repeated)["iterations"][0]["tool_calls"] + + assert [call["tool_name"] for call in tool_calls] == ["search", "search"] + assert [call["tool_input"] for call in tool_calls] == [{"q": "second"}, {"q": "second"}] + assert [call["tool_output"] for call in tool_calls] == ["second result", "second result"] + assert [call["time_cost"] for call in tool_calls] == [2, 2] + + +def test_single_call_renders_one_entry() -> None: + thought = _agent_thought( + tool="search", + tool_input=json.dumps({"search": {"q": "only"}}), + observation=json.dumps({"search": "only result"}), + tool_meta_str=json.dumps({"search": {"time_cost": 1, "tool_config": {"tool_provider_type": "builtin"}}}), + ) + + tool_calls = _get_agent_logs(thought)["iterations"][0]["tool_calls"] + + assert len(tool_calls) == 1 + assert tool_calls[0]["tool_input"] == {"q": "only"} + assert tool_calls[0]["tool_output"] == "only result" + assert tool_calls[0]["time_cost"] == 1 + + +def test_distinct_tools_render_their_own_payloads() -> None: + thought = _agent_thought( + tool="search;calculator", + tool_input=json.dumps({"search": {"q": "a"}, "calculator": {"expr": "1+1"}}), + observation=json.dumps({"search": "search result", "calculator": "2"}), + tool_meta_str=json.dumps({"search": {"time_cost": 1}, "calculator": {"time_cost": 2}}), + ) + + tool_calls = _get_agent_logs(thought)["iterations"][0]["tool_calls"] + + assert [call["tool_name"] for call in tool_calls] == ["search", "calculator"] + assert [call["tool_input"] for call in tool_calls] == [{"q": "a"}, {"expr": "1+1"}] + assert [call["tool_output"] for call in tool_calls] == ["search result", "2"]