diff --git a/camel/agents/chat_agent.py b/camel/agents/chat_agent.py index 3e8e89af08..b8c39273c3 100644 --- a/camel/agents/chat_agent.py +++ b/camel/agents/chat_agent.py @@ -3763,7 +3763,11 @@ def _execute_tool( logger.warning(f"{error_msg} with result: {result}") return self._record_tool_calling( - func_name, args, result, tool_call_id, mask_output=mask_flag + func_name, + args, + result, + tool_call_id, + mask_output=mask_flag, ) async def _aexecute_tool( @@ -3883,17 +3887,77 @@ def _record_tool_calling( cast(List[MemoryRecord], func_records), ) - # Record information about this tool call + # Calculate tool cost and token usage + cost_info = self._calculate_tool_cost(assist_msg, func_msg) + # Record information about this tool call with cost tracking tool_record = ToolCallingRecord( tool_name=func_name, args=args, result=result, tool_call_id=tool_call_id, + token_usage={ + "prompt_tokens": int(cost_info["prompt_tokens"]), + "completion_tokens": int(cost_info["completion_tokens"]), + "total_tokens": int(cost_info["total_tokens"]), + }, ) self._update_last_tool_call_state(tool_record) return tool_record + def _calculate_tool_cost( + self, + assist_msg: FunctionCallingMessage, + func_msg: FunctionCallingMessage, + ) -> Dict[str, int]: + r"""Calculate the tool cost and token usage for a tool call. + + Args: + assist_msg (FunctionCallingMessage): The assistant message + as tool call input. + func_msg (FunctionCallingMessage): The function message + as tool call output. + + Returns: + Dictionary containing token usage and cost estimates. + """ + + if hasattr(self.model_backend, 'token_counter'): + try: + input_messages = assist_msg.to_openai_message( + OpenAIBackendRole.ASSISTANT + ) + output_messages = func_msg.to_openai_message( + OpenAIBackendRole.FUNCTION + ) + input_tokens = \ + self.model_backend.token_counter.count_tokens_from_messages( + [input_messages] + ) + output_tokens = \ + self.model_backend.token_counter.count_tokens_from_messages( + [output_messages] + ) + except Exception as e: + logger.error( + f"Error calculating tool call token usage tokens: {e}" + ) + input_tokens = len(assist_msg.content.split()) + output_tokens = len(func_msg.content.split()) + else: + logger.warning( + "Token counter not available. " + "Using context words count to estimate token usage." + ) + input_tokens = len(assist_msg.content.split()) + output_tokens = len(func_msg.content.split()) + + return { + "prompt_tokens": input_tokens, + "completion_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + def _stream( self, input_message: Union[BaseMessage, str], diff --git a/camel/types/agents/tool_calling_record.py b/camel/types/agents/tool_calling_record.py index 3e0b4c91e7..eb7c643763 100644 --- a/camel/types/agents/tool_calling_record.py +++ b/camel/types/agents/tool_calling_record.py @@ -11,6 +11,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. ========= +from __future__ import annotations + from typing import Any, Dict, List, Optional from pydantic import BaseModel @@ -26,6 +28,9 @@ class ToolCallingRecord(BaseModel): tool_call_id (str): The ID of the tool call, if available. images (Optional[List[str]]): List of base64-encoded images returned by the tool, if any. + token_usage (Optional[Dict[str, int]]): Token usage breakdown for this + tool call. Contains 'prompt_tokens', 'completion_tokens', and + 'total_tokens'. """ tool_name: str @@ -33,6 +38,7 @@ class ToolCallingRecord(BaseModel): result: Any tool_call_id: str images: Optional[List[str]] = None + token_usage: Optional[Dict[str, int]] = None def __str__(self) -> str: r"""Overridden version of the string function. @@ -40,13 +46,18 @@ def __str__(self) -> str: Returns: str: Modified string to represent the tool calling. """ - return ( + base_str = ( f"Tool Execution: {self.tool_name}\n" f"\tArgs: {self.args}\n" f"\tResult: {self.result}\n" ) - def as_dict(self) -> dict[str, Any]: + if self.token_usage: + base_str += f"\tToken Usage: {self.token_usage}\n" + + return base_str + + def as_dict(self) -> Dict[str, Any]: r"""Returns the tool calling record as a dictionary. Returns: diff --git a/test/agents/test_chat_agent.py b/test/agents/test_chat_agent.py index 2e6f7538eb..2b5695d5ff 100644 --- a/test/agents/test_chat_agent.py +++ b/test/agents/test_chat_agent.py @@ -1442,3 +1442,142 @@ def test_memory_setter_preserves_system_message(): assert len(new_context) > 0 assert new_context[0]['role'] == 'system' assert new_context[0]['content'] == system_content + + +def test_calculate_tool_cost(): + # Define an echo tool + def echo(text: str) -> str: + return text + + model_config = ChatGPTConfig(temperature=0, max_tokens=200, stop="") + model_backend = ModelFactory.create( + model_platform=ModelPlatformType.OPENAI, + model_type=ModelType.GPT_5_MINI, + model_config_dict=model_config.as_dict(), + ) + agent = ChatAgent( + system_message=BaseMessage.make_assistant_message( + role_name="Assistant", + content="You are a helpful assistant.", + ), + model=model_backend, + tools=[FunctionTool(echo)], + ) + + # Call agent twice + # 1st round: call tool echo(text="hello world") + tool_call_id = "call_echo_1" + first_response = ChatCompletion( + id="mock_tool_call", + choices=[ + Choice( + finish_reason="tool_calls", + index=0, + logprobs=None, + message=ChatCompletionMessage( + content=None, + refusal=None, + role="assistant", + audio=None, + function_call=None, + tool_calls=[ + ChatCompletionMessageFunctionToolCall( + id=tool_call_id, + function=Function( + arguments='{"text":"hello world"}', + name="echo", + ), + type="function", + ) + ], + ), + ) + ], + created=1, + model="gpt-5-mini", + object="chat.completion", + service_tier=None, + usage=CompletionUsage( + completion_tokens=3, prompt_tokens=3, total_tokens=6 + ), + ) + + # 2nd round: return normal assistant content without tool calls, end loop + second_response = ChatCompletion( + id="mock_final", + choices=[ + Choice( + finish_reason="stop", + index=0, + logprobs=None, + message=ChatCompletionMessage( + content="OK", + refusal=None, + role="assistant", + audio=None, + function_call=None, + tool_calls=None, + ), + ) + ], + created=2, + model="gpt-5-mini", + object="chat.completion", + service_tier=None, + usage=CompletionUsage( + completion_tokens=2, prompt_tokens=5, total_tokens=7 + ), + ) + agent.model_backend.run = MagicMock( + side_effect=[first_response, second_response] + ) + + user_msg = BaseMessage.make_user_message( + role_name="User", content="Please call echo with text 'hello world'." + ) + agent_response = agent.step(user_msg) + + tool_calls = agent_response.info["tool_calls"] + assert tool_calls and len(tool_calls) == 1 + assert tool_calls[0].tool_name == "echo" + assert tool_calls[0].token_usage is not None + + # Get expected token usage as benchmark + from camel.messages.func_message import FunctionCallingMessage + from camel.types import OpenAIBackendRole + + token_counter = agent.model_backend.token_counter + + assist_msg = FunctionCallingMessage( + role_name=agent.role_name, + role_type=agent.role_type, + meta_dict=None, + content="", + func_name="echo", + args=tool_calls[0].args, + tool_call_id=tool_calls[0].tool_call_id, + ) + func_msg = FunctionCallingMessage( + role_name=agent.role_name, + role_type=agent.role_type, + meta_dict=None, + content="", + func_name="echo", + result=tool_calls[0].result, + tool_call_id=tool_calls[0].tool_call_id, + ) + + expected_prompt = token_counter.count_tokens_from_messages( + [assist_msg.to_openai_message(OpenAIBackendRole.ASSISTANT)] + ) + expected_completion = token_counter.count_tokens_from_messages( + [func_msg.to_openai_message(OpenAIBackendRole.FUNCTION)] + ) + expected_total = expected_prompt + expected_completion + + assert tool_calls[0].token_usage.get("prompt_tokens") == expected_prompt + assert ( + tool_calls[0].token_usage.get("completion_tokens") + == expected_completion + ) + assert tool_calls[0].token_usage.get("total_tokens") == expected_total