camel-ai 0.2.14__py3-none-any.whl → 0.2.16__py3-none-any.whl

camel/agents/chat_agent.py

@@ -32,7 +32,7 @@ from typing import (
 
 from openai.types.chat import ChatCompletionMessageToolCall
 from openai.types.chat.chat_completion_message_tool_call import Function
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 
 from camel.agents.base import BaseAgent
 from camel.memories import (
@@ -59,6 +59,7 @@ from camel.types import (
 )
 from camel.utils import (
     func_string_to_callable,
+    generate_prompt_for_structured_output,
     get_model_encoding,
     get_pydantic_object_schema,
     json_to_function_code,
@@ -142,17 +143,21 @@ class ChatAgent(BaseAgent):
             (default: :obj:`None`)
         output_language (str, optional): The language to be output by the
             agent. (default: :obj:`None`)
-        tools (List[FunctionTool], optional): List of available
-            :obj:`FunctionTool`. (default: :obj:`None`)
-        external_tools (List[FunctionTool], optional): List of external tools
-            (:obj:`FunctionTool`) bind to one chat agent. When these tools
-            are called, the agent will directly return the request instead of
+        tools (Optional[List[Union[FunctionTool, Callable]]], optional): List
+            of available :obj:`FunctionTool` or :obj:`Callable`. (default:
+            :obj:`None`)
+        external_tools (Optional[List[Union[FunctionTool, Callable]]],
+            optional): List of external tools (:obj:`FunctionTool` or or
+            :obj:`Callable`) bind to one chat agent. When these tools are
+            called, the agent will directly return the request instead of
             processing it. (default: :obj:`None`)
         response_terminators (List[ResponseTerminator], optional): List of
             :obj:`ResponseTerminator` bind to one chat agent.
             (default: :obj:`None`)
         scheduling_strategy (str): name of function that defines how to select
             the next model in ModelManager. (default: :str:`round_robin`)
+        single_iteration (bool): Whether to let the agent perform only one
+            model calling at each step. (default: :obj:`False`)
     """
 
     def __init__(
@@ -165,11 +170,13 @@ class ChatAgent(BaseAgent):
         message_window_size: Optional[int] = None,
         token_limit: Optional[int] = None,
         output_language: Optional[str] = None,
-        tools: Optional[List[FunctionTool]] = None,
-        external_tools: Optional[List[FunctionTool]] = None,
+        tools: Optional[List[Union[FunctionTool, Callable]]] = None,
+        external_tools: Optional[List[Union[FunctionTool, Callable]]] = None,
         response_terminators: Optional[List[ResponseTerminator]] = None,
         scheduling_strategy: str = "round_robin",
+        single_iteration: bool = False,
     ) -> None:
+        # Initialize the system message, converting string to BaseMessage if needed
         if isinstance(system_message, str):
             system_message = BaseMessage.make_assistant_message(
                 role_name='Assistant', content=system_message
@@ -192,34 +199,35 @@ class ChatAgent(BaseAgent):
             ),
             scheduling_strategy=scheduling_strategy,
         )
-
         self.model_type = self.model_backend.model_type
 
-        # Tool registration
-        external_tools = external_tools or []
-        tools = tools or []
-        all_tools = tools + external_tools
-        self.external_tool_names = [
-            tool.get_function_name() for tool in external_tools
+        # Initialize tools
+        self.tools: List[FunctionTool] = (
+            self._initialize_tools(tools) if tools else []
+        )
+        self.external_tools: List[FunctionTool] = (
+            self._initialize_tools(external_tools) if external_tools else []
+        )
+        self.external_tool_names: List[str] = [
+            tool.get_function_name() for tool in self.external_tools
         ]
-        self.func_dict = {
-            tool.get_function_name(): tool.func for tool in all_tools
+        self.all_tools = self.tools + self.external_tools or []
+
+        # Create tool dictionaries and configure backend tools if necessary
+        self.tool_dict = {
+            tool.get_function_name(): tool for tool in self.all_tools
         }
-        self.tool_dict = {tool.get_function_name(): tool for tool in all_tools}
 
         # If the user set tools from `ChatAgent`, it will override the
         # configured tools in `BaseModelBackend`.
-        if all_tools:
+        if self.all_tools:
             logger.warning(
                 "Overriding the configured tools in `BaseModelBackend` with the tools from `ChatAgent`."
             )
             tool_schema_list = [
-                tool.get_openai_tool_schema() for tool in all_tools
+                tool.get_openai_tool_schema() for tool in self.all_tools
             ]
             self.model_backend.model_config_dict['tools'] = tool_schema_list
-            self.tool_schema_list = tool_schema_list
-
-        self.model_config_dict = self.model_backend.model_config_dict
 
         self.model_token_limit = token_limit or self.model_backend.token_limit
         context_creator = ScoreBasedContextCreator(
@@ -237,8 +245,86 @@ class ChatAgent(BaseAgent):
         self.terminated: bool = False
         self.response_terminators = response_terminators or []
         self.init_messages()
-
         self.tool_prompt_added = False
+        self.single_iteration = single_iteration
+
+    def _initialize_tools(
+        self, tools: List[Union[FunctionTool, Callable]]
+    ) -> List[FunctionTool]:
+        r"""Helper method to initialize tools as FunctionTool instances."""
+        from camel.toolkits import FunctionTool
+
+        func_tools = []
+        for tool in tools:
+            if not isinstance(tool, FunctionTool):
+                tool = FunctionTool(tool)
+            func_tools.append(tool)
+        return func_tools
+
+    def add_tool(
+        self, tool: Union[FunctionTool, Callable], is_external: bool = False
+    ) -> None:
+        r"""Add a tool to the agent, specifying if it's an external tool."""
+        # Initialize the tool
+        initialized_tool = self._initialize_tools([tool])
+
+        # Update tools or external tools based on is_external flag
+        if is_external:
+            self.external_tools = self.external_tools + initialized_tool
+            self.external_tool_names.extend(
+                tool.get_function_name() for tool in initialized_tool
+            )
+        else:
+            self.tools = self.tools + initialized_tool
+
+        # Rebuild all_tools, and tool_dict
+        self.all_tools = self.tools + self.external_tools
+        self.tool_dict = {
+            tool.get_function_name(): tool for tool in self.all_tools
+        }
+
+        tool_schema_list = [
+            tool.get_openai_tool_schema() for tool in self.all_tools
+        ]
+        self.model_backend.model_config_dict['tools'] = tool_schema_list
+
+    def remove_tool(self, tool_name: str, is_external: bool = False) -> bool:
+        r"""Remove a tool by name, specifying if it's an external tool."""
+        tool_list = self.external_tools if is_external else self.tools
+        if not tool_list:
+            return False
+
+        for tool in tool_list:
+            if tool.get_function_name() == tool_name:
+                tool_list.remove(tool)
+                if is_external:
+                    self.external_tool_names.remove(tool_name)
+                # Reinitialize the tool dictionary
+                self.all_tools = (self.tools or []) + (
+                    self.external_tools or []
+                )
+                self.tool_dict = {
+                    tool.get_function_name(): tool for tool in self.all_tools
+                }
+                tool_schema_list = [
+                    tool.get_openai_tool_schema() for tool in self.all_tools
+                ]
+                self.model_backend.model_config_dict['tools'] = (
+                    tool_schema_list
+                )
+                return True
+        return False
+
+    def list_tools(self) -> dict:
+        r"""List all tools, separated into normal and external tools."""
+        normal_tools = [
+            tool.get_function_name() for tool in (self.tools or [])
+        ]
+        external_tools = [
+            tool.get_function_name() for tool in (self.external_tools or [])
+        ]
+
+        return {"normal_tools": normal_tools, "external_tools": external_tools}
 
     # ruff: noqa: E501
     def _generate_tool_prompt(self, tool_schema_list: List[Dict]) -> str:
@@ -264,9 +350,7 @@ class ChatAgent(BaseAgent):
 
         tool_prompt_str = "\n".join(tool_prompts)
 
-        final_prompt = f'''
-    # Tool prompt
-    TOOL_PROMPT = f"""
+        final_prompt = f"""
     You have access to the following functions:
 
     {tool_prompt_str}
@@ -274,19 +358,16 @@ class ChatAgent(BaseAgent):
     If you choose to call a function ONLY reply in the following format with no
     prefix or suffix:
 
-    <function=example_function_name>{{"example_name": "example_value"}}
-    </function>
+    <function=example_function_name>{{"example_name": "example_value"}}</function>
 
     Reminder:
-    - Function calls MUST follow the specified format, start with <function=
-      and end with </function>
+    - Function calls MUST follow the specified format, start with <function= and end with </function>
     - Required parameters MUST be specified
     - Only call one function at a time
     - Put the entire function call reply on one line
     - If there is no function call available, answer the question like normal 
       with your current knowledge and do not tell the user about function calls
     """
-    '''
         return final_prompt
 
     def _parse_tool_response(self, response: str):
@@ -310,7 +391,7 @@ class ChatAgent(BaseAgent):
                 args = json.loads(args_string)
                 return {"function": function_name, "arguments": args}
             except json.JSONDecodeError as error:
-                print(f"Error parsing function arguments: {error}")
+                logger.error(f"Error parsing function arguments: {error}")
                 return None
         return None
 
@@ -342,14 +423,13 @@ class ChatAgent(BaseAgent):
         self._system_message = message
 
     def is_tools_added(self) -> bool:
-        r"""Whether OpenAI function calling is enabled for this agent.
+        r"""Whether tool calling is enabled for this agent.
 
         Returns:
-            bool: Whether OpenAI function calling is enabled for this
-                agent, determined by whether the dictionary of tools
-                is empty.
+            bool: Whether tool calling is enabled for this agent, determined
+                by whether the dictionary of tools is empty.
         """
-        return len(self.func_dict) > 0
+        return len(self.tool_dict) > 0
 
     def update_memory(
         self, message: BaseMessage, role: OpenAIBackendRole
@@ -415,7 +495,7 @@ class ChatAgent(BaseAgent):
         Args:
             session_id (str, optional): The ID of the chat session.
             usage (Dict[str, int], optional): Information about the usage of
-                the LLM model.
+                the LLM.
             termination_reasons (List[str]): The reasons for the termination
                 of the chat session.
             num_tokens (int): The number of tokens used in the chat session.
@@ -470,26 +550,23 @@ class ChatAgent(BaseAgent):
         input_message: Union[BaseMessage, str],
         response_format: Optional[Type[BaseModel]] = None,
     ) -> ChatAgentResponse:
-        r"""Performs a single step in the chat session by generating a response
+        r"""Executes a single step in the chat session, generating a response
         to the input message.
 
         Args:
-            input_message (Union[BaseMessage, str]): The input message to the
-                agent. For BaseMessage input, its `role` field that specifies
-                the role at backend may be either `user` or `assistant` but it
-                will be set to `user` anyway since for the self agent any
-                incoming message is external. For str input, the `role_name` would be `User`.
-            response_format (Optional[Type[BaseModel]], optional): A pydantic
-                model class that includes value types and field descriptions
-                used to generate a structured response by LLM. This schema
-                helps in defining the expected output format. (default:
+            input_message (Union[BaseMessage, str]): The input message for the
+                agent. If provided as a BaseMessage, the `role` is adjusted to
+                `user` to indicate an external message.
+            response_format (Optional[Type[BaseModel]], optional): A Pydantic
+                model defining the expected structure of the response. Used to
+                generate a structured response if provided. (default:
                 :obj:`None`)
 
         Returns:
-            ChatAgentResponse: A struct containing the output messages,
-                a boolean indicating whether the chat session has terminated,
-                and information about the chat session.
+            ChatAgentResponse: Contains output messages, a termination status
+                flag, and session information.
         """
+
         if (
             self.model_backend.model_config_dict.get("response_format")
             and response_format
@@ -499,173 +576,135 @@ class ChatAgent(BaseAgent):
                 "the model configuration and in the ChatAgent step."
             )
 
-        original_model_dict = self.model_backend.model_config_dict
+        self.original_model_dict = self.model_backend.model_config_dict
         if response_format and self.model_type in {"gpt-4o", "gpt-4o-mini"}:
-            self.model_backend.model_config_dict = original_model_dict.copy()
+            self.model_backend.model_config_dict = (
+                self.original_model_dict.copy()
+            )
             self.model_backend.model_config_dict["response_format"] = (
                 response_format
             )
 
+        # Convert input message to BaseMessage if necessary
         if isinstance(input_message, str):
             input_message = BaseMessage.make_user_message(
                 role_name='User', content=input_message
             )
 
-        if "llama" in self.model_type.lower():
-            if (
-                self.model_backend.model_config_dict.get("tools", None)
-                and not self.tool_prompt_added
-            ):
-                tool_prompt = self._generate_tool_prompt(self.tool_schema_list)
+        # Handle tool prompt injection if needed
+        if (
+            self.is_tools_added()
+            and not self.model_type.support_native_tool_calling
+            and not self.tool_prompt_added
+        ):
+            self._inject_tool_prompt()
 
-                tool_sys_msg = BaseMessage.make_assistant_message(
-                    role_name="Assistant",
-                    content=tool_prompt,
-                )
+        # Add user input to memory
+        self.update_memory(input_message, OpenAIBackendRole.USER)
 
-                self.update_memory(tool_sys_msg, OpenAIBackendRole.SYSTEM)
-                self.tool_prompt_added = True
+        return self._handle_step(response_format, self.single_iteration)
 
-            self.update_memory(input_message, OpenAIBackendRole.USER)
+    def _inject_tool_prompt(self) -> None:
+        r"""Generate and add the tool prompt to memory."""
+        tool_prompt = self._generate_tool_prompt(
+            self.model_backend.model_config_dict["tools"]
+        )
+        tool_msg = BaseMessage.make_assistant_message(
+            role_name="Assistant", content=tool_prompt
+        )
+        self.update_memory(tool_msg, OpenAIBackendRole.SYSTEM)
+        self.tool_prompt_added = True
 
-            tool_call_records: List[FunctionCallingRecord] = []
-            while True:
-                # Check if token has exceeded
-                try:
-                    openai_messages, num_tokens = self.memory.get_context()
-                except RuntimeError as e:
-                    return self._step_token_exceed(
-                        e.args[1], tool_call_records, "max_tokens_exceeded"
-                    )
+    def _handle_step(
+        self,
+        response_format: Optional[Type[BaseModel]],
+        single_step: bool,
+    ) -> ChatAgentResponse:
+        r"""Handles a single or multi-step interaction."""
 
-                (
-                    response,
-                    output_messages,
-                    finish_reasons,
-                    usage_dict,
-                    response_id,
-                ) = self._step_model_response(openai_messages, num_tokens)
-                # If the model response is not a function call, meaning the
-                # model has generated a message response, break the loop
-                if (
-                    not self.is_tools_added()
-                    or not isinstance(response, ChatCompletion)
-                    or "</function>" not in response.choices[0].message.content  # type: ignore[operator]
-                ):
-                    break
-
-                parsed_content = self._parse_tool_response(
-                    response.choices[0].message.content  # type: ignore[arg-type]
-                )
+        if (
+            self.model_backend.model_config_dict.get("tool_choice")
+            == "required"
+            and not single_step
+        ):
+            raise ValueError(
+                "`tool_choice` cannot be set to `required` for multi-step"
+                " mode. To proceed, set `single_iteration` to `True`."
+            )
 
-                response.choices[0].message.tool_calls = [
-                    ChatCompletionMessageToolCall(
-                        id=str(uuid.uuid4()),
-                        function=Function(
-                            arguments=str(parsed_content["arguments"]).replace(
-                                "'", '"'
-                            ),
-                            name=str(parsed_content["function"]),
-                        ),
-                        type="function",
-                    )
-                ]
+        # Record function calls made during the session
+        tool_call_records: List[FunctionCallingRecord] = []
 
-                # Check for external tool call
-                tool_call_request = response.choices[0].message.tool_calls[0]
-                if tool_call_request.function.name in self.external_tool_names:
-                    # if model calls an external tool, directly return the
-                    # request
-                    info = self._step_get_info(
-                        output_messages,
-                        finish_reasons,
-                        usage_dict,
-                        response_id,
-                        tool_call_records,
-                        num_tokens,
-                        tool_call_request,
-                    )
-                    return ChatAgentResponse(
-                        msgs=output_messages,
-                        terminated=self.terminated,
-                        info=info,
-                    )
+        external_tool_request = None
 
-                # Normal function calling
-                tool_call_records.append(
-                    self._step_tool_call_and_update(response)
+        while True:
+            try:
+                openai_messages, num_tokens = self.memory.get_context()
+            except RuntimeError as e:
+                self.model_backend.model_config_dict = self.original_model_dict
+                return self._step_token_exceed(
+                    e.args[1], tool_call_records, "max_tokens_exceeded"
                 )
 
-            if response_format is not None:
-                (
-                    output_messages,
-                    finish_reasons,
-                    usage_dict,
-                    response_id,
-                    tool_call,
-                    num_tokens,
-                ) = self._structure_output_with_function(response_format)
-                tool_call_records.append(tool_call)
+            # Prompt engineering approach for structured output for non-native tool calling models
+            inject_prompt_for_structured_output = (
+                response_format
+                and not self.model_type.support_native_structured_output
+            )
+
+            if inject_prompt_for_structured_output:
+                # update last openai message
+                usr_msg = openai_messages.pop()
+                usr_msg["content"] = generate_prompt_for_structured_output(
+                    response_format,
+                    usr_msg["content"],  # type: ignore [arg-type]
+                )
+                openai_messages.append(usr_msg)
 
-            info = self._step_get_info(
+            # Process model response
+            (
+                response,
                 output_messages,
                 finish_reasons,
                 usage_dict,
                 response_id,
-                tool_call_records,
-                num_tokens,
-            )
-
-            if len(output_messages) == 1:
-                # Auto record if the output result is a single message
-                self.record_message(output_messages[0])
-            else:
-                logger.warning(
-                    "Multiple messages returned in `step()`, message won't be "
-                    "recorded automatically. Please call `record_message()` "
-                    "to record the selected message manually."
-                )
-
-            return ChatAgentResponse(
-                msgs=output_messages, terminated=self.terminated, info=info
-            )
-
-        else:
-            self.update_memory(input_message, OpenAIBackendRole.USER)
+            ) = self._step_model_response(openai_messages, num_tokens)
 
-            tool_call_records: List[FunctionCallingRecord] = []  # type: ignore[no-redef]
-            while True:
-                # Check if token has exceeded
+            # Try to parse structured output to return a Pydantic object
+            if inject_prompt_for_structured_output and isinstance(
+                response, ChatCompletion
+            ):
+                content = response.choices[0].message.content
                 try:
-                    openai_messages, num_tokens = self.memory.get_context()
-                except RuntimeError as e:
-                    self.model_backend.model_config_dict = original_model_dict
-                    return self._step_token_exceed(
-                        e.args[1], tool_call_records, "max_tokens_exceeded"
+                    json_content = json.loads(str(content))
+                    output_messages[0].parsed = response_format(**json_content)  # type: ignore [assignment, misc]
+                except json.JSONDecodeError as e:
+                    logger.error(
+                        f"Failed in parsing the output into JSON: {e}"
+                    )
+                    output_messages[0].parsed = None
+                except ValidationError as e:
+                    logger.warning(
+                        "Successfully generating JSON response, "
+                        "but failed in parsing it into Pydantic object :"
+                        f"{e}, return the JSON response in parsed field"
                     )
+                    output_messages[0].parsed = json_content
 
-                (
-                    response,
-                    output_messages,
-                    finish_reasons,
-                    usage_dict,
-                    response_id,
-                ) = self._step_model_response(openai_messages, num_tokens)
-                # If the model response is not a function call, meaning the
-                # model has generated a message response, break the loop
-                if (
-                    not self.is_tools_added()
-                    or not isinstance(response, ChatCompletion)
-                    or not response.choices[0].message.tool_calls
-                ):
-                    break
-
-                # Check for external tool call
-                tool_call_request = response.choices[0].message.tool_calls[0]
-
-                if tool_call_request.function.name in self.external_tool_names:
-                    # if model calls an external tool, directly return the
-                    # request
+            # Finalize on standard response in multi-step mode
+            if self._is_standard_response(response):
+                break
+
+            # Handle tool requests
+            tool_request = self._extract_tool_call(response)
+            if isinstance(response, ChatCompletion) and tool_request:
+                response.choices[0].message.tool_calls = [tool_request]
+                tool_call_records.append(
+                    self._step_tool_call_and_update(response)
+                )
+
+                if tool_request.function.name in self.external_tool_names:
+                    external_tool_request = tool_request
                     info = self._step_get_info(
                         output_messages,
                         finish_reasons,
@@ -673,58 +712,132 @@ class ChatAgent(BaseAgent):
                         response_id,
                         tool_call_records,
                         num_tokens,
-                        tool_call_request,
+                        tool_request,
+                    )
+                    self._log_final_output(output_messages)
+                    self.model_backend.model_config_dict = (
+                        self.original_model_dict
                     )
-
-                    self.model_backend.model_config_dict = original_model_dict
                     return ChatAgentResponse(
                         msgs=output_messages,
                         terminated=self.terminated,
                         info=info,
                     )
 
-                # Normal function calling
-                tool_call_records.append(
-                    self._step_tool_call_and_update(response)
-                )
-
-            if (
-                response_format is not None
-                and self.model_type.support_native_tool_calling
-                and self.model_type not in {"gpt-4o", "gpt-4o-mini"}
-            ):
-                (
-                    output_messages,
-                    finish_reasons,
-                    usage_dict,
-                    response_id,
-                    tool_call,
-                    num_tokens,
-                ) = self._structure_output_with_function(response_format)
-                tool_call_records.append(tool_call)
+            # Single-step mode ends after one iteration
+            if single_step:
+                break
 
-            info = self._step_get_info(
+        # Optional structured output via function calling
+        if (
+            response_format
+            and not inject_prompt_for_structured_output
+            and self.model_type
+            not in {
+                "gpt-4o",
+                "gpt-4o-mini",
+            }
+        ):
+            (
                 output_messages,
                 finish_reasons,
                 usage_dict,
                 response_id,
-                tool_call_records,
+                tool_call,
                 num_tokens,
-            )
+            ) = self._structure_output_with_function(response_format)
+            tool_call_records.append(tool_call)
+
+        # Final info and response
+        info = self._step_get_info(
+            output_messages,
+            finish_reasons,
+            usage_dict,
+            response_id,
+            tool_call_records,
+            num_tokens,
+            external_tool_request,
+        )
+        self._log_final_output(output_messages)
+        self.model_backend.model_config_dict = self.original_model_dict
+        return ChatAgentResponse(
+            msgs=output_messages, terminated=self.terminated, info=info
+        )
+
+    def _extract_tool_call(
+        self, response: Any
+    ) -> Optional[ChatCompletionMessageToolCall]:
+        r"""Extract the tool call from the model response, if present.
+
+        Args:
+            response (Any): The model's response object.
 
-            if len(output_messages) == 1:
-                # Auto record if the output result is a single message
-                self.record_message(output_messages[0])
-            else:
-                logger.warning(
-                    "Multiple messages returned in `step()`, message won't be "
-                    "recorded automatically. Please call `record_message()` "
-                    "to record the selected message manually."
+        Returns:
+            Optional[ChatCompletionMessageToolCall]: The parsed tool call if
+                present, otherwise None.
+        """
+        # Check if the response contains tool calls
+        if (
+            self.is_tools_added()
+            and not self.model_type.support_native_tool_calling
+            and "</function>" in response.choices[0].message.content
+        ):
+            parsed_content = self._parse_tool_response(
+                response.choices[0].message.content
+            )
+            if parsed_content:
+                return ChatCompletionMessageToolCall(
+                    id=str(uuid.uuid4()),
+                    function=Function(
+                        arguments=str(parsed_content["arguments"]).replace(
+                            "'", '"'
+                        ),
+                        name=str(parsed_content["function"]),
+                    ),
+                    type="function",
                 )
+        elif (
+            self.is_tools_added()
+            and self.model_type.support_native_tool_calling
+            and response.choices[0].message.tool_calls
+        ):
+            return response.choices[0].message.tool_calls[0]
+
+        # No tool call found
+        return None
+
+    def _is_standard_response(self, response: Any) -> bool:
+        r"""Determine if the provided response is a standard reply without
+        tool calls.
+
+        Args:
+            response (Any): The response object to evaluate.
+
+        Returns:
+            bool: `True` if the response is a standard reply, `False`
+                otherwise.
+        """
+        if not self.is_tools_added():
+            return True
 
-            self.model_backend.model_config_dict = original_model_dict
-            return ChatAgentResponse(
-                msgs=output_messages, terminated=self.terminated, info=info
+        if not isinstance(response, ChatCompletion):
+            return True
+
+        if self.model_type.support_native_tool_calling:
+            return response.choices[0].message.tool_calls is None
+
+        return "</function>" not in str(
+            response.choices[0].message.content or ""
+        )
+
+    def _log_final_output(self, output_messages: List[BaseMessage]) -> None:
+        r"""Log final messages or warnings about multiple responses."""
+        if len(output_messages) == 1:
+            self.record_message(output_messages[0])
+        else:
+            logger.warning(
+                "Multiple messages returned in `step()`. Record "
+                "selected message manually using `record_message()`."
             )
 
     async def step_async(
@@ -740,7 +853,8 @@ class ChatAgent(BaseAgent):
                 agent. For BaseMessage input, its `role` field that specifies
                 the role at backend may be either `user` or `assistant` but it
                 will be set to `user` anyway since for the self agent any
-                incoming message is external. For str input, the `role_name` would be `User`.
+                incoming message is external. For str input, the `role_name`
+                would be `User`.
             response_format (Optional[Type[BaseModel]], optional): A pydantic
                 model class that includes value types and field descriptions
                 used to generate a structured response by LLM. This schema
@@ -779,13 +893,13 @@ class ChatAgent(BaseAgent):
             if (
                 not self.is_tools_added()
                 or not isinstance(response, ChatCompletion)
-                or response.choices[0].message.tool_calls is None
+                or not response.choices[0].message.tool_calls
             ):
                 break
 
             # Check for external tool call
-            tool_call_request = response.choices[0].message.tool_calls[0]
-            if tool_call_request.function.name in self.external_tool_names:
+            external_tool_request = response.choices[0].message.tool_calls[0]
+            if external_tool_request.function.name in self.external_tool_names:
                 # if model calls an external tool, directly return the request
                 info = self._step_get_info(
                     output_messages,
@@ -794,7 +908,7 @@ class ChatAgent(BaseAgent):
                     response_id,
                     tool_call_records,
                     num_tokens,
-                    tool_call_request,
+                    external_tool_request,
                 )
                 return ChatAgentResponse(
                     msgs=output_messages, terminated=self.terminated, info=info
@@ -859,7 +973,7 @@ class ChatAgent(BaseAgent):
 
         # Perform function calling
         func_assistant_msg, func_result_msg, tool_call_record = (
-            self.step_tool_call(response)
+            self._step_tool_call(response)
         )
 
         # Update the messages
@@ -913,11 +1027,9 @@ class ChatAgent(BaseAgent):
         func_callable = func_string_to_callable(func_str)
         func = FunctionTool(func_callable)
 
-        original_func_dict = self.func_dict
         original_model_dict = self.model_backend.model_config_dict
 
         # Replace the original tools with the structuring function
-        self.func_dict = {func.get_function_name(): func.func}
         self.tool_dict = {func.get_function_name(): func}
         self.model_backend.model_config_dict = original_model_dict.copy()
         self.model_backend.model_config_dict["tools"] = [
@@ -945,7 +1057,6 @@ class ChatAgent(BaseAgent):
             base_message_item.content = json.dumps(tool_call_record.result)
 
         # Recover the original tools
-        self.func_dict = original_func_dict
         self.model_backend.model_config_dict = original_model_dict
 
         return (
@@ -1245,7 +1356,7 @@ class ChatAgent(BaseAgent):
             info=info,
         )
 
-    def step_tool_call(
+    def _step_tool_call(
         self,
         response: ChatCompletion,
     ) -> Tuple[
@@ -1268,7 +1379,9 @@ class ChatAgent(BaseAgent):
             raise RuntimeError("Tool call is None")
         func_name = choice.message.tool_calls[0].function.name
 
-        args = json.loads(choice.message.tool_calls[0].function.arguments)
+        arguments_str = choice.message.tool_calls[0].function.arguments
+        args = self._safe_json_loads(arguments_str)
+
         tool = self.tool_dict[func_name]
         result = tool(**args)
 
@@ -1295,6 +1408,18 @@ class ChatAgent(BaseAgent):
         )
         return assist_msg, func_msg, func_record
 
+    def _safe_json_loads(self, arguments_str):
+        # Replace Python types with their JSON equivalents
+        arguments_str = arguments_str.replace("None", "null")
+        arguments_str = arguments_str.replace("True", "true")
+        arguments_str = arguments_str.replace("False", "false")
+
+        # Attempt to parse the corrected string
+        try:
+            return json.loads(arguments_str)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON format: {e}")
+
     async def step_tool_call_async(
         self,
         response: ChatCompletion,