From e6ac4008fe202320942b4c6e0713ee0495868d7d Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 Apr 2026 15:35:50 +0100 Subject: [PATCH] feat: better tool names for llms --- application/agents/base.py | 4 +- application/agents/tool_executor.py | 99 ++++++++---- .../agents/tools/tool_action_parser.py | 31 +++- application/api/v1/translator.py | 24 +-- application/llm/handlers/base.py | 17 +- tests/agents/test_base_agent.py | 6 +- tests/agents/test_tool_action_parser.py | 66 ++++++++ tests/agents/test_tool_executor.py | 150 ++++++++++++++---- tests/llm/handlers/test_llm_handlers.py | 1 + tests/llm/test_base.py | 10 +- tests/test_agent_token_tracking.py | 2 + tests/test_client_tools.py | 29 ++-- tests/test_remaining_coverage.py | 3 +- tests/test_v1_translator.py | 54 +++---- 14 files changed, 353 insertions(+), 143 deletions(-) diff --git a/application/agents/base.py b/application/agents/base.py index 4576e98d..596caa5a 100644 --- a/application/agents/base.py +++ b/application/agents/base.py @@ -218,7 +218,7 @@ class BaseAgent(ABC): "data": { "tool_name": pending.get("tool_name", "unknown"), "call_id": call_id, - "action_name": f"{pending['action_name']}_{pending['tool_id']}", + "action_name": pending.get("llm_name", pending["name"]), "arguments": args, "status": "denied", }, @@ -242,7 +242,7 @@ class BaseAgent(ABC): "data": { "tool_name": pending.get("tool_name", "unknown"), "call_id": call_id, - "action_name": f"{pending['action_name']}_{pending['tool_id']}", + "action_name": pending.get("llm_name", pending["name"]), "arguments": args, "result": ( result_str[:50] + "..." diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py index dcfe5f3f..11095b10 100644 --- a/application/agents/tool_executor.py +++ b/application/agents/tool_executor.py @@ -1,6 +1,7 @@ import logging import uuid -from typing import Dict, List, Optional +from collections import Counter +from typing import Dict, List, Optional, Tuple from bson.objectid import ObjectId @@ -32,6 +33,8 @@ class ToolExecutor: self._loaded_tools: Dict[str, object] = {} self.conversation_id: Optional[str] = None self.client_tools: Optional[List[Dict]] = None + self._name_to_tool: Dict[str, Tuple[str, str]] = {} + self._tool_to_name: Dict[Tuple[str, str], str] = {} def get_tools(self) -> Dict[str, Dict]: """Load tool configs from DB based on user context. @@ -98,8 +101,6 @@ class ToolExecutor: for i, ct in enumerate(client_tools): func = ct.get("function", ct) # tolerate bare {"name":..} too name = func.get("name", f"clienttool{i}") - # Tool IDs must NOT contain underscores — ToolActionParser splits - # on "_" and takes the last segment as the ID. tool_id = f"ct{i}" tools_dict[tool_id] = { @@ -117,8 +118,21 @@ class ToolExecutor: return tools_dict def prepare_tools_for_llm(self, tools_dict: Dict) -> List[Dict]: - """Convert tool configs to LLM function schemas.""" - result = [] + """Convert tool configs to LLM function schemas. + + Action names are kept clean for the LLM: + - Unique action names appear as-is (e.g. ``get_weather``). + - Duplicate action names get numbered suffixes (e.g. ``search_1``, + ``search_2``). + + A reverse mapping is stored in ``_name_to_tool`` so that tool calls + can be routed back to the correct ``(tool_id, action_name)`` without + brittle string splitting. + """ + # Pass 1: collect entries and count action name occurrences + entries: List[Tuple[str, str, Dict, bool]] = [] # (tool_id, action_name, action, is_client) + name_counts: Counter = Counter() + for tool_id, tool in tools_dict.items(): is_api = tool["name"] == "api_tool" is_client = tool.get("client_side", False) @@ -137,21 +151,48 @@ class ToolExecutor: for action in actions: if not action.get("active", True): continue + entries.append((tool_id, action["name"], action, is_client)) + name_counts[action["name"]] += 1 - # Client-side tools already have parameters in standard format - if is_client: - params = action.get("parameters", {}) - else: - params = self._build_tool_parameters(action) + # Pass 2: assign LLM-visible names and build mappings + self._name_to_tool = {} + self._tool_to_name = {} + collision_counters: Dict[str, int] = {} + all_llm_names: set = set() - result.append({ - "type": "function", - "function": { - "name": f"{action['name']}_{tool_id}", - "description": action.get("description", ""), - "parameters": params, - }, - }) + result = [] + for tool_id, action_name, action, is_client in entries: + if name_counts[action_name] == 1: + llm_name = action_name + else: + counter = collision_counters.get(action_name, 1) + candidate = f"{action_name}_{counter}" + # Skip if candidate collides with a unique action name + while candidate in all_llm_names or ( + candidate in name_counts and name_counts[candidate] == 1 + ): + counter += 1 + candidate = f"{action_name}_{counter}" + collision_counters[action_name] = counter + 1 + llm_name = candidate + + all_llm_names.add(llm_name) + self._name_to_tool[llm_name] = (tool_id, action_name) + self._tool_to_name[(tool_id, action_name)] = llm_name + + if is_client: + params = action.get("parameters", {}) + else: + params = self._build_tool_parameters(action) + + result.append({ + "type": "function", + "function": { + "name": llm_name, + "description": action.get("description", ""), + "parameters": params, + }, + }) return result def _build_tool_parameters(self, action: Dict) -> Dict: @@ -176,9 +217,10 @@ class ToolExecutor: Returns a dict describing the pending action if pause is needed, None otherwise. """ - parser = ToolActionParser(llm_class_name) + parser = ToolActionParser(llm_class_name, name_mapping=self._name_to_tool) tool_id, action_name, call_args = parser.parse_args(call) call_id = getattr(call, "id", None) or str(uuid.uuid4()) + llm_name = getattr(call, "name", "") if tool_id is None or action_name is None or tool_id not in tools_dict: return None # Will be handled as error by execute() @@ -189,10 +231,11 @@ class ToolExecutor: if tool_data.get("client_side"): return { "call_id": call_id, - "name": getattr(call, "name", f"{action_name}_{tool_id}"), + "name": llm_name, "tool_name": tool_data.get("name", "unknown"), "tool_id": tool_id, "action_name": action_name, + "llm_name": llm_name, "arguments": call_args if isinstance(call_args, dict) else {}, "pause_type": "requires_client_execution", "thought_signature": getattr(call, "thought_signature", None), @@ -212,10 +255,11 @@ class ToolExecutor: if action_data.get("require_approval"): return { "call_id": call_id, - "name": getattr(call, "name", f"{action_name}_{tool_id}"), + "name": llm_name, "tool_name": tool_data.get("name", "unknown"), "tool_id": tool_id, "action_name": action_name, + "llm_name": llm_name, "arguments": call_args if isinstance(call_args, dict) else {}, "pause_type": "awaiting_approval", "thought_signature": getattr(call, "thought_signature", None), @@ -225,21 +269,22 @@ class ToolExecutor: def execute(self, tools_dict: Dict, call, llm_class_name: str): """Execute a tool call. Yields status events, returns (result, call_id).""" - parser = ToolActionParser(llm_class_name) + parser = ToolActionParser(llm_class_name, name_mapping=self._name_to_tool) tool_id, action_name, call_args = parser.parse_args(call) + llm_name = getattr(call, "name", "unknown") call_id = getattr(call, "id", None) or str(uuid.uuid4()) if tool_id is None or action_name is None: - error_message = f"Error: Failed to parse LLM tool call. Tool name: {getattr(call, 'name', 'unknown')}" + error_message = f"Error: Failed to parse LLM tool call. Tool name: {llm_name}" logger.error(error_message) tool_call_data = { "tool_name": "unknown", "call_id": call_id, - "action_name": getattr(call, "name", "unknown"), + "action_name": llm_name, "arguments": call_args or {}, - "result": f"Failed to parse tool call. Invalid tool name format: {getattr(call, 'name', 'unknown')}", + "result": f"Failed to parse tool call. Invalid tool name format: {llm_name}", } yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}} self.tool_calls.append(tool_call_data) @@ -252,7 +297,7 @@ class ToolExecutor: tool_call_data = { "tool_name": "unknown", "call_id": call_id, - "action_name": f"{action_name}_{tool_id}", + "action_name": llm_name, "arguments": call_args, "result": f"Tool with ID {tool_id} not found. Available tools: {list(tools_dict.keys())}", } @@ -263,7 +308,7 @@ class ToolExecutor: tool_call_data = { "tool_name": tools_dict[tool_id]["name"], "call_id": call_id, - "action_name": f"{action_name}_{tool_id}", + "action_name": llm_name, "arguments": call_args, } yield {"type": "tool_call", "data": {**tool_call_data, "status": "pending"}} diff --git a/application/agents/tools/tool_action_parser.py b/application/agents/tools/tool_action_parser.py index aff7d2e2..2c31ad45 100644 --- a/application/agents/tools/tool_action_parser.py +++ b/application/agents/tools/tool_action_parser.py @@ -5,8 +5,9 @@ logger = logging.getLogger(__name__) class ToolActionParser: - def __init__(self, llm_type): + def __init__(self, llm_type, name_mapping=None): self.llm_type = llm_type + self.name_mapping = name_mapping self.parsers = { "OpenAILLM": self._parse_openai_llm, "GoogleLLM": self._parse_google_llm, @@ -16,22 +17,33 @@ class ToolActionParser: parser = self.parsers.get(self.llm_type, self._parse_openai_llm) return parser(call) + def _resolve_via_mapping(self, call_name): + """Look up (tool_id, action_name) from the name mapping if available.""" + if self.name_mapping and call_name in self.name_mapping: + return self.name_mapping[call_name] + return None + def _parse_openai_llm(self, call): try: call_args = json.loads(call.arguments) + + resolved = self._resolve_via_mapping(call.name) + if resolved: + return resolved[0], resolved[1], call_args + + # Fallback: legacy split on "_" for backward compatibility tool_parts = call.name.split("_") - # If the tool name doesn't contain an underscore, it's likely a hallucinated tool if len(tool_parts) < 2: logger.warning( - f"Invalid tool name format: {call.name}. Expected format: action_name_tool_id" + f"Invalid tool name format: {call.name}. " + "Could not resolve via mapping or legacy parsing." ) return None, None, None tool_id = tool_parts[-1] action_name = "_".join(tool_parts[:-1]) - # Validate that tool_id looks like a numerical ID if not tool_id.isdigit(): logger.warning( f"Tool ID '{tool_id}' is not numerical. This might be a hallucinated tool call." @@ -45,19 +57,24 @@ class ToolActionParser: def _parse_google_llm(self, call): try: call_args = call.arguments + + resolved = self._resolve_via_mapping(call.name) + if resolved: + return resolved[0], resolved[1], call_args + + # Fallback: legacy split on "_" for backward compatibility tool_parts = call.name.split("_") - # If the tool name doesn't contain an underscore, it's likely a hallucinated tool if len(tool_parts) < 2: logger.warning( - f"Invalid tool name format: {call.name}. Expected format: action_name_tool_id" + f"Invalid tool name format: {call.name}. " + "Could not resolve via mapping or legacy parsing." ) return None, None, None tool_id = tool_parts[-1] action_name = "_".join(tool_parts[:-1]) - # Validate that tool_id looks like a numerical ID if not tool_id.isdigit(): logger.warning( f"Tool ID '{tool_id}' is not numerical. This might be a hallucinated tool call." diff --git a/application/api/v1/translator.py b/application/api/v1/translator.py index 9212bea3..3d3c64fe 100644 --- a/application/api/v1/translator.py +++ b/application/api/v1/translator.py @@ -7,21 +7,17 @@ This module handles: """ import json -import re import time from typing import Any, Dict, List, Optional -# Pattern matching internal tool-id suffixes (e.g. _ct0, _ct12) -_TOOL_SUFFIX_RE = re.compile(r"_ct\d+$") +def _get_client_tool_name(tc: Dict) -> str: + """Return the original tool name for client-facing responses. - -def _strip_tool_suffix(name: str) -> str: - """Remove internal tool-id suffix from a tool name for client responses. - - Internally tools are named ``action_ct0`` so the LLM can route calls. - Standard OpenAI clients expect the original registered name back. + For client-side tools the ``tool_name`` field carries the name the + client originally registered. Fall back to ``action_name`` (which + is now the clean LLM-visible name) or ``name``. """ - return _TOOL_SUFFIX_RE.sub("", name) + return tc.get("tool_name", tc.get("action_name", tc.get("name", ""))) # --------------------------------------------------------------------------- @@ -214,9 +210,7 @@ def translate_response( "id": tc.get("call_id", ""), "type": "function", "function": { - "name": _strip_tool_suffix( - tc.get("action_name", tc.get("name", "")) - ), + "name": _get_client_tool_name(tc), "arguments": ( json.dumps(tc["arguments"]) if isinstance(tc.get("arguments"), dict) @@ -358,9 +352,7 @@ def translate_stream_event( "id": tc_data.get("call_id", ""), "type": "function", "function": { - "name": _strip_tool_suffix( - tc_data.get("action_name", "") - ), + "name": _get_client_tool_name(tc_data), "arguments": args_str, }, }], diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index 815ce179..a2edd233 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -803,7 +803,7 @@ class LLMHandler(ABC): "data": { "tool_name": pause_info["tool_name"], "call_id": pause_info["call_id"], - "action_name": f"{pause_info['action_name']}_{pause_info['tool_id']}", + "action_name": pause_info.get("llm_name", pause_info["name"]), "arguments": pause_info["arguments"], "status": pause_info["pause_type"], }, @@ -857,16 +857,15 @@ class LLMHandler(ABC): error_message = self.create_tool_message(error_call, error_response) updated_messages.append(error_message) - call_parts = call.name.split("_") - if len(call_parts) >= 2: - tool_id = call_parts[-1] # Last part is tool ID (e.g., "1") - action_name = "_".join(call_parts[:-1]) - tool_name = tools_dict.get(tool_id, {}).get("name", "unknown_tool") - full_action_name = f"{action_name}_{tool_id}" + mapping = agent.tool_executor._name_to_tool + if call.name in mapping: + resolved_tool_id, _ = mapping[call.name] + tool_name = tools_dict.get(resolved_tool_id, {}).get( + "name", "unknown_tool" + ) else: tool_name = "unknown_tool" - action_name = call.name - full_action_name = call.name + full_action_name = call.name yield { "type": "tool_call", "data": { diff --git a/tests/agents/test_base_agent.py b/tests/agents/test_base_agent.py index 05e831f0..f65c7b0d 100644 --- a/tests/agents/test_base_agent.py +++ b/tests/agents/test_base_agent.py @@ -341,7 +341,7 @@ class TestBaseAgentTools: assert len(agent.tools) == 1 assert agent.tools[0]["type"] == "function" - assert agent.tools[0]["function"]["name"] == "get_data_1" + assert agent.tools[0]["function"]["name"] == "get_data" def test_prepare_tools_with_regular_tool( self, agent_base_params, mock_llm_creator, mock_llm_handler_creator @@ -365,7 +365,7 @@ class TestBaseAgentTools: agent._prepare_tools(tools_dict) assert len(agent.tools) == 1 - assert agent.tools[0]["function"]["name"] == "action1_1" + assert agent.tools[0]["function"]["name"] == "action1" def test_prepare_tools_filters_inactive_actions( self, agent_base_params, mock_llm_creator, mock_llm_handler_creator @@ -395,7 +395,7 @@ class TestBaseAgentTools: agent._prepare_tools(tools_dict) assert len(agent.tools) == 1 - assert agent.tools[0]["function"]["name"] == "active_action_1" + assert agent.tools[0]["function"]["name"] == "active_action" @pytest.mark.unit diff --git a/tests/agents/test_tool_action_parser.py b/tests/agents/test_tool_action_parser.py index 840d070c..61e9e962 100644 --- a/tests/agents/test_tool_action_parser.py +++ b/tests/agents/test_tool_action_parser.py @@ -202,3 +202,69 @@ class TestToolActionParser: assert action_name == "create_record" assert call_args["data"]["name"] == "John" assert call_args["data"]["age"] == 30 + + +@pytest.mark.unit +class TestToolActionParserWithMapping: + """Tests for the mapping-based lookup path.""" + + def test_openai_mapping_resolves_clean_name(self): + mapping = {"get_weather": ("ct0", "get_weather")} + parser = ToolActionParser("OpenAILLM", name_mapping=mapping) + + call = Mock() + call.name = "get_weather" + call.arguments = '{"city": "SF"}' + + tool_id, action_name, call_args = parser.parse_args(call) + assert tool_id == "ct0" + assert action_name == "get_weather" + assert call_args == {"city": "SF"} + + def test_openai_mapping_resolves_numbered_suffix(self): + mapping = {"search_1": ("t1", "search"), "search_2": ("t2", "search")} + parser = ToolActionParser("OpenAILLM", name_mapping=mapping) + + call = Mock() + call.name = "search_1" + call.arguments = '{"q": "test"}' + + tool_id, action_name, call_args = parser.parse_args(call) + assert tool_id == "t1" + assert action_name == "search" + + def test_google_mapping_resolves(self): + mapping = {"get_weather": ("ct0", "get_weather")} + parser = ToolActionParser("GoogleLLM", name_mapping=mapping) + + call = Mock() + call.name = "get_weather" + call.arguments = {"city": "SF"} + + tool_id, action_name, call_args = parser.parse_args(call) + assert tool_id == "ct0" + assert action_name == "get_weather" + + def test_fallback_to_split_when_not_in_mapping(self): + mapping = {"get_weather": ("ct0", "get_weather")} + parser = ToolActionParser("OpenAILLM", name_mapping=mapping) + + call = Mock() + call.name = "unknown_action_99" + call.arguments = "{}" + + tool_id, action_name, call_args = parser.parse_args(call) + # Falls back to legacy split + assert tool_id == "99" + assert action_name == "unknown_action" + + def test_no_mapping_uses_legacy_split(self): + parser = ToolActionParser("OpenAILLM", name_mapping=None) + + call = Mock() + call.name = "action_123" + call.arguments = '{"k": "v"}' + + tool_id, action_name, call_args = parser.parse_args(call) + assert tool_id == "123" + assert action_name == "action" diff --git a/tests/agents/test_tool_executor.py b/tests/agents/test_tool_executor.py index bb3038fe..02b6ce0f 100644 --- a/tests/agents/test_tool_executor.py +++ b/tests/agents/test_tool_executor.py @@ -80,7 +80,7 @@ class TestToolExecutorPrepare: result = executor.prepare_tools_for_llm(tools_dict) assert len(result) == 1 assert result[0]["type"] == "function" - assert result[0]["function"]["name"] == "do_thing_t1" + assert result[0]["function"]["name"] == "do_thing" assert "query" in result[0]["function"]["parameters"]["properties"] def test_prepare_tools_skips_inactive_actions(self): @@ -97,7 +97,96 @@ class TestToolExecutorPrepare: result = executor.prepare_tools_for_llm(tools_dict) assert len(result) == 1 - assert result[0]["function"]["name"] == "active_one_t1" + assert result[0]["function"]["name"] == "active_one" + + def test_prepare_tools_builds_name_mapping(self): + executor = ToolExecutor() + tools_dict = { + "t1": { + "name": "test_tool", + "actions": [ + {"name": "do_thing", "description": "D", "active": True, "parameters": {"properties": {}}}, + ], + } + } + executor.prepare_tools_for_llm(tools_dict) + assert executor._name_to_tool["do_thing"] == ("t1", "do_thing") + assert executor._tool_to_name[("t1", "do_thing")] == "do_thing" + + def test_prepare_tools_duplicate_names_get_numbered_suffixes(self): + executor = ToolExecutor() + tools_dict = { + "t1": { + "name": "tool_a", + "actions": [ + {"name": "search", "description": "D", "active": True, "parameters": {"properties": {}}}, + ], + }, + "t2": { + "name": "tool_b", + "actions": [ + {"name": "search", "description": "D", "active": True, "parameters": {"properties": {}}}, + ], + }, + } + result = executor.prepare_tools_for_llm(tools_dict) + names = [r["function"]["name"] for r in result] + assert "search_1" in names + assert "search_2" in names + assert executor._name_to_tool["search_1"][1] == "search" + assert executor._name_to_tool["search_2"][1] == "search" + + def test_prepare_tools_unique_name_no_suffix(self): + executor = ToolExecutor() + tools_dict = { + "t1": { + "name": "tool_a", + "actions": [ + {"name": "get_weather", "description": "D", "active": True, "parameters": {"properties": {}}}, + ], + }, + "t2": { + "name": "tool_b", + "actions": [ + {"name": "send_email", "description": "D", "active": True, "parameters": {"properties": {}}}, + ], + }, + } + result = executor.prepare_tools_for_llm(tools_dict) + names = [r["function"]["name"] for r in result] + assert "get_weather" in names + assert "send_email" in names + + def test_prepare_tools_suffix_skips_collision_with_unique_name(self): + """If action 'foo_1' exists as unique and 'foo' is duplicated, skip '_1'.""" + executor = ToolExecutor() + tools_dict = { + "t1": { + "name": "tool_a", + "actions": [ + {"name": "foo", "description": "D", "active": True, "parameters": {"properties": {}}}, + ], + }, + "t2": { + "name": "tool_b", + "actions": [ + {"name": "foo", "description": "D", "active": True, "parameters": {"properties": {}}}, + ], + }, + "t3": { + "name": "tool_c", + "actions": [ + {"name": "foo_1", "description": "D", "active": True, "parameters": {"properties": {}}}, + ], + }, + } + result = executor.prepare_tools_for_llm(tools_dict) + names = [r["function"]["name"] for r in result] + # foo_1 is taken by the unique action, so duplicates skip to _2 and _3 + assert "foo_1" in names # The unique action + assert "foo_2" in names + assert "foo_3" in names + assert executor._name_to_tool["foo_1"] == ("t3", "foo_1") def test_build_tool_parameters_filters_non_llm_fields(self): executor = ToolExecutor() @@ -139,17 +228,10 @@ class TestCheckPause: call.thought_signature = None return call - def test_client_side_tool_returns_suffixed_name(self, monkeypatch): - """check_pause returns the LLM-facing suffixed name for internal routing.""" + def test_client_side_tool_returns_llm_name(self): + """check_pause returns the clean LLM-facing name and llm_name field.""" executor = ToolExecutor() - monkeypatch.setattr( - "application.agents.tool_executor.ToolActionParser", - lambda _cls: Mock( - parse_args=Mock(return_value=("ct0", "write_file", {"path": "test.md"})) - ), - ) - tools_dict = { "ct0": { "name": "write_file", @@ -160,27 +242,22 @@ class TestCheckPause: } } - call = self._make_call(name="write_file_ct0") - result = executor.check_pause(tools_dict, call, "MockLLM") + # Prepare tools so the mapping is built + executor.prepare_tools_for_llm(tools_dict) + + call = self._make_call(name="write_file") + result = executor.check_pause(tools_dict, call, "OpenAILLM") assert result is not None - # name keeps the suffix for LLM message reconstruction during continuation - assert result["name"] == "write_file_ct0" - # action_name is the clean parsed name + assert result["name"] == "write_file" + assert result["llm_name"] == "write_file" assert result["action_name"] == "write_file" assert result["tool_id"] == "ct0" - def test_approval_required_returns_suffixed_name(self, monkeypatch): - """check_pause for approval-required tools also returns suffixed name.""" + def test_approval_required_returns_llm_name(self): + """check_pause for approval-required tools returns clean LLM name.""" executor = ToolExecutor() - monkeypatch.setattr( - "application.agents.tool_executor.ToolActionParser", - lambda _cls: Mock( - parse_args=Mock(return_value=("t1", "delete_all", {})) - ), - ) - tools_dict = { "t1": { "name": "dangerous_tool", @@ -191,11 +268,14 @@ class TestCheckPause: } } - call = self._make_call(name="delete_all_t1") - result = executor.check_pause(tools_dict, call, "MockLLM") + executor.prepare_tools_for_llm(tools_dict) + + call = self._make_call(name="delete_all") + result = executor.check_pause(tools_dict, call, "OpenAILLM") assert result is not None - assert result["name"] == "delete_all_t1" + assert result["name"] == "delete_all" + assert result["llm_name"] == "delete_all" assert result["action_name"] == "delete_all" @@ -214,7 +294,7 @@ class TestToolExecutorExecute: monkeypatch.setattr( "application.agents.tool_executor.ToolActionParser", - lambda _cls: Mock(parse_args=Mock(return_value=(None, None, {}))), + lambda _cls, **kw: Mock(parse_args=Mock(return_value=(None, None, {}))), ) call = self._make_call(name="bad") @@ -238,7 +318,7 @@ class TestToolExecutorExecute: monkeypatch.setattr( "application.agents.tool_executor.ToolActionParser", - lambda _cls: Mock(parse_args=Mock(return_value=("missing_id", "action", {}))), + lambda _cls, **kw: Mock(parse_args=Mock(return_value=("missing_id", "action", {}))), ) call = self._make_call() @@ -261,7 +341,7 @@ class TestToolExecutorExecute: monkeypatch.setattr( "application.agents.tool_executor.ToolActionParser", - lambda _cls: Mock(parse_args=Mock(return_value=("t1", "test_action", {"param1": "val"}))), + lambda _cls, **kw: Mock(parse_args=Mock(return_value=("t1", "test_action", {"param1": "val"}))), ) tools_dict = { @@ -315,7 +395,7 @@ class TestToolExecutorExecute: monkeypatch.setattr( "application.agents.tool_executor.ToolActionParser", - lambda _cls: Mock(parse_args=Mock(return_value=("t1", "test_action", {}))), + lambda _cls, **kw: Mock(parse_args=Mock(return_value=("t1", "test_action", {}))), ) tools_dict = { @@ -355,7 +435,7 @@ class TestToolExecutorExecute: monkeypatch.setattr( "application.agents.tool_executor.ToolActionParser", - lambda _cls: Mock( + lambda _cls, **kw: Mock( parse_args=Mock(return_value=("t1", "get_users", {"body_param": "val"})) ), ) @@ -402,7 +482,7 @@ class TestToolExecutorExecute: monkeypatch.setattr( "application.agents.tool_executor.ToolActionParser", - lambda _cls: Mock( + lambda _cls, **kw: Mock( parse_args=Mock(return_value=("t1", "act", {})) ), ) @@ -447,7 +527,7 @@ class TestToolExecutorExecute: monkeypatch.setattr( "application.agents.tool_executor.ToolActionParser", - lambda _cls: Mock( + lambda _cls, **kw: Mock( parse_args=Mock(return_value=("t1", "act", {"q": "v"})) ), ) diff --git a/tests/llm/handlers/test_llm_handlers.py b/tests/llm/handlers/test_llm_handlers.py index da693ec2..6baf04e8 100644 --- a/tests/llm/handlers/test_llm_handlers.py +++ b/tests/llm/handlers/test_llm_handlers.py @@ -678,6 +678,7 @@ class TestHandleToolCalls: agent.context_limit_reached = False agent.llm.__class__.__name__ = "MockLLM" agent.tool_executor.check_pause = Mock(return_value=None) + agent.tool_executor._name_to_tool = {} agent._execute_tool_action = Mock(side_effect=RuntimeError("exec error")) call = ToolCall(id="c1", name="action_1", arguments="{}") diff --git a/tests/llm/test_base.py b/tests/llm/test_base.py index 613e541a..e62d0186 100644 --- a/tests/llm/test_base.py +++ b/tests/llm/test_base.py @@ -480,11 +480,12 @@ class TestHandleToolCallsErrors: agent._check_context_limit = MagicMock(return_value=False) agent.llm.__class__.__name__ = "MockLLM" agent.tool_executor.check_pause = MagicMock(return_value=None) + agent.tool_executor._name_to_tool = {"search": ("1", "search")} agent._execute_tool_action = MagicMock( side_effect=RuntimeError("tool failed") ) - tool_call = ToolCall(id="tc1", name="search_1", arguments={"q": "test"}) + tool_call = ToolCall(id="tc1", name="search", arguments={"q": "test"}) tools_dict = {"1": {"name": "search_tool"}} messages = [{"role": "user", "content": "hi"}] @@ -510,6 +511,7 @@ class TestHandleToolCallsErrors: agent._check_context_limit = MagicMock(return_value=False) agent.llm.__class__.__name__ = "MockLLM" agent.tool_executor.check_pause = MagicMock(return_value=None) + agent.tool_executor._name_to_tool = {} agent._execute_tool_action = MagicMock( side_effect=RuntimeError("tool failed") ) @@ -1175,12 +1177,13 @@ class TestHandleToolCallsErrorsAdditional: agent._check_context_limit = MagicMock(return_value=False) agent.llm.__class__.__name__ = "MockLLM" agent.tool_executor.check_pause = MagicMock(return_value=None) + agent.tool_executor._name_to_tool = {"do_thing": ("42", "do_thing")} agent._execute_tool_action = MagicMock( side_effect=RuntimeError("broken tool") ) tool_call = ToolCall( - id="tc1", name="do_thing_42", arguments={"x": 1} + id="tc1", name="do_thing", arguments={"x": 1} ) tools_dict = {"42": {"name": "my_tool"}} messages = [{"role": "user", "content": "go"}] @@ -1211,7 +1214,7 @@ class TestHandleToolCallsErrorsAdditional: ] assert len(error_events) == 1 assert error_events[0]["data"]["tool_name"] == "my_tool" - assert error_events[0]["data"]["action_name"] == "do_thing_42" + assert error_events[0]["data"]["action_name"] == "do_thing" def test_tool_error_with_no_context_check(self): """Cover line 660: messages.copy() at start of handle_tool_calls.""" @@ -1221,6 +1224,7 @@ class TestHandleToolCallsErrorsAdditional: agent.llm.__class__.__name__ = "MockLLM" agent.tool_executor = MagicMock() agent.tool_executor.check_pause = MagicMock(return_value=None) + agent.tool_executor._name_to_tool = {} agent._execute_tool_action = MagicMock( side_effect=ValueError("bad args") ) diff --git a/tests/test_agent_token_tracking.py b/tests/test_agent_token_tracking.py index b1ab1a99..bdb21549 100644 --- a/tests/test_agent_token_tracking.py +++ b/tests/test_agent_token_tracking.py @@ -178,6 +178,7 @@ class TestLLMHandlerTokenTracking: mock_agent.context_limit_reached = False mock_agent.llm.__class__.__name__ = "MockLLM" mock_agent.tool_executor.check_pause = Mock(return_value=None) + mock_agent.tool_executor._name_to_tool = {} call_count = [0] @@ -239,6 +240,7 @@ class TestLLMHandlerTokenTracking: mock_agent._check_context_limit = Mock(return_value=False) mock_agent.llm.__class__.__name__ = "MockLLM" mock_agent.tool_executor.check_pause = Mock(return_value=None) + mock_agent.tool_executor._name_to_tool = {} mock_agent._execute_tool_action = Mock( return_value=iter([{"type": "tool_call", "data": {}}]) ) diff --git a/tests/test_client_tools.py b/tests/test_client_tools.py index fa8d434b..6c4cd6d8 100644 --- a/tests/test_client_tools.py +++ b/tests/test_client_tools.py @@ -142,7 +142,7 @@ class TestPrepareClientToolsForLlm: assert len(schemas) == 1 assert schemas[0]["type"] == "function" - assert schemas[0]["function"]["name"] == "get_weather_ct0" + assert schemas[0]["function"]["name"] == "get_weather" assert schemas[0]["function"]["description"] == "Get weather" # Parameters passed through directly (not filtered by _build_tool_parameters) assert "city" in schemas[0]["function"]["parameters"]["properties"] @@ -187,8 +187,8 @@ class TestPrepareClientToolsForLlm: assert len(schemas) == 2 names = {s["function"]["name"] for s in schemas} - assert "do_thing_t1" in names - assert "local_fn_ct0" in names + assert "do_thing" in names + assert "local_fn" in names # --------------------------------------------------------------------------- @@ -251,7 +251,8 @@ class TestCheckPauseClientTools: ], } } - call = self._make_call(name="get_weather_ct0") + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(name="get_weather") result = executor.check_pause(tools_dict, call, "OpenAILLM") assert result is not None @@ -269,7 +270,8 @@ class TestCheckPauseClientTools: ], } } - call = self._make_call(name="search_0") + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(name="search") result = executor.check_pause(tools_dict, call, "OpenAILLM") assert result is None @@ -315,19 +317,21 @@ class TestHandlerClientToolPause: # check_pause returns pause info for client tool agent.tool_executor.check_pause = Mock(return_value={ "call_id": "c1", - "name": "get_weather_ct0", + "name": "get_weather", "tool_name": "get_weather", "tool_id": "ct0", "action_name": "get_weather", + "llm_name": "get_weather", "arguments": {"city": "SF"}, "pause_type": "requires_client_execution", "thought_signature": None, }) + agent.tool_executor._name_to_tool = {"get_weather": ("ct0", "get_weather")} # Simulate streaming: one chunk with tool_calls finish_reason chunk = LLMResponse( content="", - tool_calls=[ToolCall(id="c1", name="get_weather_ct0", arguments='{"city": "SF"}', index=0)], + tool_calls=[ToolCall(id="c1", name="get_weather", arguments='{"city": "SF"}', index=0)], finish_reason="tool_calls", raw_response={}, ) @@ -371,10 +375,11 @@ class TestHandlerClientToolPause: if call_count["n"] == 2: # Second tool is client-side return { "call_id": "c2", - "name": "get_weather_ct0", + "name": "get_weather", "tool_name": "get_weather", "tool_id": "ct0", "action_name": "get_weather", + "llm_name": "get_weather", "arguments": {}, "pause_type": "requires_client_execution", "thought_signature": None, @@ -382,6 +387,10 @@ class TestHandlerClientToolPause: return None agent.tool_executor.check_pause = Mock(side_effect=check_pause_fn) + agent.tool_executor._name_to_tool = { + "search": ("0", "search"), + "get_weather": ("ct0", "get_weather"), + } def fake_execute(tools_dict, call): yield {"type": "tool_call", "data": {"status": "pending"}} @@ -390,8 +399,8 @@ class TestHandlerClientToolPause: agent._execute_tool_action = Mock(side_effect=fake_execute) calls = [ - ToolCall(id="c1", name="search_0", arguments="{}"), - ToolCall(id="c2", name="get_weather_ct0", arguments="{}"), + ToolCall(id="c1", name="search", arguments="{}"), + ToolCall(id="c2", name="get_weather", arguments="{}"), ] gen = handler.handle_tool_calls( diff --git a/tests/test_remaining_coverage.py b/tests/test_remaining_coverage.py index e27f86f2..62d430bb 100644 --- a/tests/test_remaining_coverage.py +++ b/tests/test_remaining_coverage.py @@ -1238,7 +1238,8 @@ class TestEmbeddingPipelineAddDocWithRetry: # NUL characters should be removed assert "\x00" not in doc.page_content - def test_add_text_to_store_with_retry_failure(self): + @patch("time.sleep", return_value=None) + def test_add_text_to_store_with_retry_failure(self, _mock_sleep): from application.parser.embedding_pipeline import add_text_to_store_with_retry mock_store = MagicMock() diff --git a/tests/test_v1_translator.py b/tests/test_v1_translator.py index 1f33e31c..92deb912 100644 --- a/tests/test_v1_translator.py +++ b/tests/test_v1_translator.py @@ -9,7 +9,7 @@ import json import pytest from application.api.v1.translator import ( - _strip_tool_suffix, + _get_client_tool_name, convert_history, extract_tool_results, is_continuation, @@ -297,13 +297,13 @@ class TestTranslateResponse: ) assert resp["docsgpt"]["tool_calls"] == tool_calls - def test_pending_tool_calls_strips_ct_suffix(self): - """Internal _ct\\d+ suffixes must be stripped from tool names in responses.""" + def test_pending_tool_calls_uses_tool_name(self): + """Client tool responses use the original tool_name, not the LLM-visible action_name.""" pending = [ { "call_id": "c1", - "name": "get_weather_ct0", - "action_name": "get_weather_ct0", + "tool_name": "get_weather", + "action_name": "get_weather", "arguments": {"city": "SF"}, } ] @@ -319,13 +319,13 @@ class TestTranslateResponse: tc = resp["choices"][0]["message"]["tool_calls"][0] assert tc["function"]["name"] == "get_weather" - def test_pending_tool_calls_non_ct_suffix_preserved(self): - """Non-client tool suffixes (e.g. _t1) should not be stripped.""" + def test_pending_tool_calls_tool_name_takes_precedence(self): + """When tool_name differs from action_name, tool_name is used.""" pending = [ { "call_id": "c1", - "name": "search_t1", - "action_name": "search_t1", + "tool_name": "search", + "action_name": "search_1", "arguments": {"q": "test"}, } ] @@ -339,8 +339,7 @@ class TestTranslateResponse: pending_tool_calls=pending, ) tc = resp["choices"][0]["message"]["tool_calls"][0] - # _t1 is NOT a client-tool suffix (_ct\d+), so it stays - assert tc["function"]["name"] == "search_t1" + assert tc["function"]["name"] == "search" def test_pending_tool_calls(self): pending = [ @@ -445,14 +444,15 @@ class TestTranslateStreamEvent: assert tc["id"] == "c1" assert tc["function"]["name"] == "get_weather" - def test_tool_call_client_execution_strips_ct_suffix(self): - """Internal _ct suffixes must be stripped from streaming tool call names.""" + def test_tool_call_client_execution_uses_tool_name(self): + """Streaming tool calls use tool_name (original name) for client responses.""" chunks = translate_stream_event( { "type": "tool_call", "data": { "call_id": "c1", - "action_name": "create_ct0", + "tool_name": "create", + "action_name": "create", "arguments": {"title": "test"}, "status": "requires_client_execution", }, @@ -557,27 +557,21 @@ class TestTranslateStreamEvent: # --------------------------------------------------------------------------- -# _strip_tool_suffix +# _get_client_tool_name # --------------------------------------------------------------------------- @pytest.mark.unit -class TestStripToolSuffix: +class TestGetClientToolName: - def test_strips_ct0(self): - assert _strip_tool_suffix("create_ct0") == "create" + def test_uses_tool_name_when_present(self): + assert _get_client_tool_name({"tool_name": "create", "action_name": "create_1"}) == "create" - def test_strips_ct_multi_digit(self): - assert _strip_tool_suffix("write_file_ct12") == "write_file" + def test_falls_back_to_action_name(self): + assert _get_client_tool_name({"action_name": "get_weather"}) == "get_weather" - def test_preserves_non_ct_suffix(self): - assert _strip_tool_suffix("search_t1") == "search_t1" + def test_falls_back_to_name(self): + assert _get_client_tool_name({"name": "search"}) == "search" - def test_preserves_plain_name(self): - assert _strip_tool_suffix("get_weather") == "get_weather" - - def test_preserves_empty(self): - assert _strip_tool_suffix("") == "" - - def test_ct_in_middle_not_stripped(self): - assert _strip_tool_suffix("ct0_action") == "ct0_action" + def test_returns_empty_when_no_fields(self): + assert _get_client_tool_name({}) == ""