langroid 0.6.7__tar.gz → 0.8.0__tar.gz

{langroid-0.6.7 → langroid-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langroid
-Version: 0.6.7
+Version: 0.8.0
 Summary: Harness LLMs with Multi-Agent Programming
 License: MIT
 Author: Prasad Chalasani
@@ -235,6 +235,8 @@ teacher_task.run()
 <details>
 <summary> <b>Click to expand</b></summary>
 
+- **Aug 2024:** 
+  - **[0.7.0](https://github.com/langroid/langroid/releases/tag/0.7.0)** OpenAI tools API support.
 - **Jul 2024:**
   - **[0.3.0](https://github.com/langroid/langroid/releases/tag/0.3.0)**: Added [FastEmbed](https://qdrant.github.io/fastembed/qdrant/Usage_With_Qdrant/) embeddings from Qdrant
 - **Jun 2024:**

{langroid-0.6.7 → langroid-0.8.0}/README.md

@@ -129,6 +129,8 @@ teacher_task.run()
 <details>
 <summary> <b>Click to expand</b></summary>
 
+- **Aug 2024:** 
+  - **[0.7.0](https://github.com/langroid/langroid/releases/tag/0.7.0)** OpenAI tools API support.
 - **Jul 2024:**
   - **[0.3.0](https://github.com/langroid/langroid/releases/tag/0.3.0)**: Added [FastEmbed](https://qdrant.github.io/fastembed/qdrant/Usage_With_Qdrant/) embeddings from Qdrant
 - **Jun 2024:**

{langroid-0.6.7 → langroid-0.8.0}/langroid/agent/base.py

@@ -1,4 +1,5 @@
 import asyncio
+import copy
 import inspect
 import json
 import logging
@@ -33,6 +34,7 @@ from langroid.language_models.base import (
     LLMMessage,
     LLMResponse,
     LLMTokenUsage,
+    OpenAIToolCall,
     StreamingIfAllowed,
 )
 from langroid.language_models.openai_gpt import OpenAIGPT, OpenAIGPTConfig
@@ -93,6 +95,11 @@ class Agent(ABC):
     """
 
     id: str = Field(default_factory=lambda: ObjectRegistry.new_id())
+    # OpenAI tool-calls awaiting response; update when a tool result with Role.TOOL
+    # is added to self.message_history
+    oai_tool_calls: List[OpenAIToolCall] = []
+    # Index of ALL tool calls generated by the agent
+    oai_tool_id2call: Dict[str, OpenAIToolCall] = {}
 
     def __init__(self, config: AgentConfig = AgentConfig()):
         self.config = config
@@ -343,30 +350,129 @@ class Agent(ABC):
             )
             return results
         if not settings.quiet:
+            results_str = (
+                results if isinstance(results, str) else json.dumps(results, indent=2)
+            )
             console.print(f"[red]{self.indent}", end="")
-            print(f"[red]Agent: {escape(results)}")
-            maybe_json = len(extract_top_level_json(results)) > 0
+            print(f"[red]Agent: {escape(results_str)}")
+            maybe_json = len(extract_top_level_json(results_str)) > 0
             self.callbacks.show_agent_response(
-                content=results,
+                content=results_str,
                 language="json" if maybe_json else "text",
             )
         sender_name = self.config.name
         if isinstance(msg, ChatDocument) and msg.function_call is not None:
             # if result was from handling an LLM `function_call`,
-            # set sender_name to "request", i.e. name of the function_call
+            # set sender_name to name of the function_call
             sender_name = msg.function_call.name
 
+        results_str, id2result, oai_tool_id = self._process_tool_results(
+            results if isinstance(results, str) else "",
+            id2result=None if isinstance(results, str) else results,
+            tool_calls=(msg.oai_tool_calls if isinstance(msg, ChatDocument) else None),
+        )
         return ChatDocument(
-            content=results,
+            content=results_str,
+            oai_tool_id2result=id2result,
             metadata=ChatDocMetaData(
                 source=Entity.AGENT,
                 sender=Entity.AGENT,
                 sender_name=sender_name,
+                oai_tool_id=oai_tool_id,
                 # preserve trail of tool_ids for OpenAI Assistant fn-calls
                 tool_ids=[] if isinstance(msg, str) else msg.metadata.tool_ids,
             ),
         )
 
+    def _process_tool_results(
+        self,
+        results: str,
+        id2result: Dict[str, str] | None,
+        tool_calls: List[OpenAIToolCall] | None = None,
+    ) -> Tuple[str, Dict[str, str] | None, str | None]:
+        """
+        Process results from a response, based on whether
+        they are results of OpenAI tool-calls from THIS agent.
+
+        Return:
+            - str: The response string
+            - Dict[str,str]|None: A dict of OpenAI tool id -> result, if there are
+                multiple tool results.
+            - str|None: tool_id if we there was a single tool result
+
+        """
+        id2result_ = copy.deepcopy(id2result) if id2result is not None else None
+        results_str = ""
+        oai_tool_id = None
+
+        if results != "":
+            # in this case ignore id2result
+            assert (
+                id2result is None
+            ), "id2result should be None when results is not empty!"
+            results_str = results
+            if len(self.oai_tool_calls) > 0:
+                # There may be multiple tool-calls, but we only have one result,
+                # so we return it as a map of the first tool-call id to the result.
+                assert (
+                    len(self.oai_tool_calls) == 1
+                ), "There are multiple pending tool-calls, but only one result!"
+                # We record the tool_id of the tool-call that
+                # the result is a response to, so that ChatDocument.to_LLMMessage
+                # can properly set the `tool_call_id` field of the LLMMessage.
+                oai_tool_id = self.oai_tool_calls[0].id
+        elif id2result is not None and id2result_ is not None:  # appease mypy
+            assert (
+                tool_calls is not None
+            ), "tool_calls cannot be None when id2result is not None!"
+            # This must be an OpenAI tool id -> result map;
+            # However some ids may not correspond to the tool-calls in the list of
+            # pending tool-calls (self.oai_tool_calls). Such results are concatenated
+            # into a simple string, to store in the ChatDocument.content,
+            # and the rest
+            # (i.e. those that DO correspond to tools in self.oai_tool_calls)
+            # are stored as a dict in ChatDocument.oa_tool_id2result.
+
+            # OAI tools from THIS agent, awaiting response
+            pending_tool_ids = [tc.id for tc in self.oai_tool_calls]
+            # tool_calls that the results are a response to
+            # (but these may have been sent from another agent, hence may not be in
+            # self.oai_tool_calls)
+            parent_tool_id2name = {
+                tc.id: tc.function.name
+                for tc in tool_calls or []
+                if tc.function is not None
+            }
+
+            # (id, result) for result NOT corresponding to self.oai_tool_calls,
+            # i.e. these are results of EXTERNAL tool-calls from another agent.
+            external_tool_id_results = []
+
+            for tc_id, result in id2result.items():
+                if tc_id not in pending_tool_ids:
+                    external_tool_id_results.append((tc_id, result))
+                    id2result_.pop(tc_id)
+            if len(external_tool_id_results) == 0:
+                results_str = ""
+            elif len(external_tool_id_results) == 1:
+                results_str = external_tool_id_results[0][1]
+            else:
+                results_str = "\n\n".join(
+                    [
+                        f"Result from tool/function {parent_tool_id2name[id]}: {result}"
+                        for id, result in external_tool_id_results
+                    ]
+                )
+
+            if len(id2result_) == 0:
+                id2result_ = None
+            elif len(id2result_) == 1 and len(external_tool_id_results) == 0:
+                results_str = list(id2result_.values())[0]
+                oai_tool_id = list(id2result_.keys())[0]
+                id2result_ = None
+
+        return results_str, id2result_, oai_tool_id
+
     def _response_template(self, e: Entity, content: str | None = None) -> ChatDocument:
         """Template for response from entity `e`."""
         return ChatDocument(
@@ -627,16 +733,25 @@ class Agent(ABC):
             # (either via OpenAI Fn-call or Langroid-native ToolMessage)
             return msg.tool_messages
         assert isinstance(msg, ChatDocument)
-        # when `content` is non-empty, we assume there will be no `function_call`
-        if msg.content != "":
+        if (
+            msg.content != ""
+            and msg.oai_tool_calls is None
+            and msg.function_call is None
+        ):
+
             tools = self.get_json_tool_messages(msg.content)
             msg.tool_messages = tools
             return tools
 
-        # otherwise, we look for a `function_call`
-        fun_call_cls = self.get_function_call_class(msg)
-        tools = [fun_call_cls] if fun_call_cls is not None else []
+        # otherwise, we look for `tool_calls` (possibly multiple)
+        tools = self.get_oai_tool_calls_classes(msg)
         msg.tool_messages = tools
+
+        if len(tools) == 0:
+            # otherwise, we look for a `function_call`
+            fun_call_cls = self.get_function_call_class(msg)
+            tools = [fun_call_cls] if fun_call_cls is not None else []
+            msg.tool_messages = tools
         return tools
 
     def get_json_tool_messages(self, input_str: str) -> List[ToolMessage]:
@@ -656,6 +771,10 @@ class Agent(ABC):
         return [r for r in results if r is not None]
 
     def get_function_call_class(self, msg: ChatDocument) -> Optional[ToolMessage]:
+        """
+        From ChatDocument (constructed from an LLM Response), get the `ToolMessage`
+        corresponding to the `function_call` if it exists.
+        """
         if msg.function_call is None:
             return None
         tool_name = msg.function_call.name
@@ -677,6 +796,39 @@ class Agent(ABC):
         tool = tool_class.parse_obj(tool_msg)
         return tool
 
+    def get_oai_tool_calls_classes(self, msg: ChatDocument) -> List[ToolMessage]:
+        """
+        From ChatDocument (constructed from an LLM Response), get
+         a list of ToolMessages corresponding to the `tool_calls`, if any.
+        """
+
+        if msg.oai_tool_calls is None:
+            return []
+        tools = []
+        for tc in msg.oai_tool_calls:
+            if tc.function is None:
+                continue
+            tool_name = tc.function.name
+            tool_msg = tc.function.arguments or {}
+            if tool_name not in self.llm_tools_handled:
+                logger.warning(
+                    f"""
+                    The tool_call '{tool_name}' is not handled 
+                    by the agent named '{self.config.name}'!
+                    If you intended this agent to handle this function_call,
+                    either the fn-call name is incorrectly generated by the LLM,
+                    (in which case you may need to adjust your LLM instructions),
+                    or you need to enable this agent to handle this fn-call.
+                    """
+                )
+                continue
+            tool_class = self.llm_tools_map[tool_name]
+            tool_msg.update(dict(request=tool_name))
+            tool = tool_class.parse_obj(tool_msg)
+            tool.id = tc.id or ""
+            tools.append(tool)
+        return tools
+
     def tool_validation_error(self, ve: ValidationError) -> str:
         """
         Handle a validation error raised when parsing a tool message,
@@ -699,7 +851,9 @@ class Agent(ABC):
         Please write your message again, correcting the errors.
         """
 
-    def handle_message(self, msg: str | ChatDocument) -> None | str | ChatDocument:
+    def handle_message(
+        self, msg: str | ChatDocument
+    ) -> None | str | Dict[str, str] | ChatDocument:
         """
         Handle a "tool" message either a string containing one or more
         valid "tool" JSON substrings,  or a
@@ -711,9 +865,16 @@ class Agent(ABC):
             msg (str | ChatDocument): The string or ChatDocument to handle
 
         Returns:
-            Optional[Str]: The result of the handler method in string form so it can
-            be sent back to the LLM, or None if `msg` was not successfully
-            handled by a method.
+            The result of the handler method can be:
+             - None if no tools successfully handled, or no tools present
+             - str if langroid-native JSON tools were handled, and results concatenated,
+                 OR there's a SINGLE OpenAI tool-call.
+                (We do this so the common scenario of a single tool/fn-call
+                 has a simple behavior).
+             - Dict[str, str] if multiple OpenAI tool-calls were handled
+                 (dict is an id->result map)
+             - ChatDocument if a handler returned a ChatDocument, intended to be the
+                 final response of the `agent_response` method.
         """
         try:
             tools = self.get_tool_messages(msg)
@@ -729,14 +890,25 @@ class Agent(ABC):
             return None
         if len(tools) == 0:
             return self.handle_message_fallback(msg)
-
+        has_ids = all([t.id != "" for t in tools])
         results = [self.handle_tool_message(t) for t in tools]
-
-        results_list = [r for r in results if r is not None]
-        if len(results_list) == 0:
+        tool_names = [t.default_value("request") for t in tools]
+        if has_ids:
+            id2result = {
+                t.id: r
+                for t, r in zip(tools, results)
+                if r is not None and isinstance(r, str)
+            }
+
+        name_results_list = [
+            (name, r) for name, r in zip(tool_names, results) if r is not None
+        ]
+        if len(name_results_list) == 0:
             return None  # self.handle_message_fallback(msg)
         # there was a non-None result
-        chat_doc_results = [r for r in results_list if isinstance(r, ChatDocument)]
+        chat_doc_results = [
+            r for _, r in name_results_list if isinstance(r, ChatDocument)
+        ]
         if len(chat_doc_results) > 1:
             logger.warning(
                 """There were multiple ChatDocument results from tools,
@@ -747,8 +919,21 @@ class Agent(ABC):
         if len(chat_doc_results) > 0:
             return chat_doc_results[0]
 
-        str_doc_results = [r for r in results_list if isinstance(r, str)]
-        final = "\n".join(str_doc_results)
+        if has_ids and len(id2result) > 1:
+            # if there are multiple OpenAI Tool results, return them as a dict
+            return id2result
+
+        if len(name_results_list) == 1 and isinstance(name_results_list[0][1], str):
+            # single str result -- return it as is
+            return name_results_list[0][1]
+
+        # multi-results: prepend the tool name to each result
+        str_results = [
+            f"Result from {name}: {r}"
+            for name, r in name_results_list
+            if isinstance(r, str)
+        ]
+        final = "\n\n".join(str_results)
         return final
 
     def handle_message_fallback(

{langroid-0.6.7 → langroid-0.8.0}/langroid/agent/callbacks/chainlit.py

@@ -58,7 +58,7 @@ async def setup_llm() -> None:
     timeout = llm_settings.get("timeout", 90)
     logger.info(f"Using model: {model}")
     llm_config = lm.OpenAIGPTConfig(
-        chat_model=model or lm.OpenAIChatModel.GPT4_TURBO,
+        chat_model=model or lm.OpenAIChatModel.GPT4o,
         # or, other possibilities for example:
         # "litellm/ollama_chat/mistral"
         # "litellm/ollama_chat/mistral:7b-instruct-v0.2-q8_0"

{langroid-0.6.7 → langroid-0.8.0}/langroid/agent/chat_agent.py

@@ -16,8 +16,10 @@ from langroid.language_models.base import (
     LLMFunctionSpec,
     LLMMessage,
     LLMResponse,
+    OpenAIToolSpec,
     Role,
     StreamingIfAllowed,
+    ToolChoiceTypes,
 )
 from langroid.language_models.openai_gpt import OpenAIGPT
 from langroid.utils.configuration import settings
@@ -39,14 +41,19 @@ class ChatAgentConfig(AgentConfig):
         user_message: user message to include in message sequence.
              Used only if `task` is not specified in the constructor.
         use_tools: whether to use our own ToolMessages mechanism
-        use_functions_api: whether to use functions native to the LLM API
-                (e.g. OpenAI's `function_call` mechanism)
+        use_functions_api: whether to use functions/tools native to the LLM API
+                (e.g. OpenAI's `function_call` or `tool_call` mechanism)
+        use_tools_api: When `use_functions_api` is True, if this is also True,
+            the OpenAI tool-call API is used, rather than the older/deprecated
+            function-call API. However the tool-call API has some tricky aspects,
+            hence we set this to False by default.
     """
 
     system_message: str = "You are a helpful assistant."
     user_message: Optional[str] = None
     use_tools: bool = False
     use_functions_api: bool = True
+    use_tools_api: bool = False
 
     def _set_fn_or_tools(self, fn_available: bool) -> None:
         """
@@ -205,6 +212,23 @@ class ChatAgent(Agent):
             msgs.append(LLMMessage(role=Role.USER, content=self.user_message))
         return msgs
 
+    def _drop_msg_update_tool_calls(self, msg: LLMMessage) -> None:
+        id2idx = {t.id: i for i, t in enumerate(self.oai_tool_calls)}
+        if msg.role == Role.TOOL:
+            # dropping tool result, so ADD the corresponding tool-call back
+            # to the list of pending calls!
+            id = msg.tool_call_id
+            if id in self.oai_tool_id2call:
+                self.oai_tool_calls.append(self.oai_tool_id2call[id])
+        elif msg.tool_calls is not None:
+            # dropping a msg with tool-calls, so DROP these from pending list
+            # as well as from id -> call map
+            for tool_call in msg.tool_calls:
+                if tool_call.id in id2idx:
+                    self.oai_tool_calls.pop(id2idx[tool_call.id])
+                if tool_call.id in self.oai_tool_id2call:
+                    del self.oai_tool_id2call[tool_call.id]
+
     def clear_history(self, start: int = -2) -> None:
         """
         Clear the message history, starting at the index `start`
@@ -218,7 +242,10 @@ class ChatAgent(Agent):
             n = len(self.message_history)
             start = max(0, n + start)
         dropped = self.message_history[start:]
-        for msg in dropped:
+        # consider the dropped msgs in REVERSE order, so we are
+        # carefully updating self.oai_tool_calls
+        for msg in reversed(dropped):
+            self._drop_msg_update_tool_calls(msg)
             # clear out the chat document from the ObjectRegistry
             ChatDocument.delete_id(msg.chat_document_id)
         self.message_history = self.message_history[:start]
@@ -519,9 +546,14 @@ class ChatAgent(Agent):
         hist, output_len = self._prep_llm_messages(message)
         if len(hist) == 0:
             return None
+        tool_choice = (
+            "auto"
+            if isinstance(message, str)
+            else (message.oai_tool_choice if message is not None else "auto")
+        )
         with StreamingIfAllowed(self.llm, self.llm.get_stream()):
-            response = self.llm_response_messages(hist, output_len)
-        self.message_history.append(ChatDocument.to_LLMMessage(response))
+            response = self.llm_response_messages(hist, output_len, tool_choice)
+        self.message_history.extend(ChatDocument.to_LLMMessage(response))
         response.metadata.msg_idx = len(self.message_history) - 1
         response.metadata.agent_id = self.id
         # Preserve trail of tool_ids for OpenAI Assistant fn-calls
@@ -543,9 +575,16 @@ class ChatAgent(Agent):
         hist, output_len = self._prep_llm_messages(message)
         if len(hist) == 0:
             return None
+        tool_choice = (
+            "auto"
+            if isinstance(message, str)
+            else (message.oai_tool_choice if message is not None else "auto")
+        )
         with StreamingIfAllowed(self.llm, self.llm.get_stream()):
-            response = await self.llm_response_messages_async(hist, output_len)
-        self.message_history.append(ChatDocument.to_LLMMessage(response))
+            response = await self.llm_response_messages_async(
+                hist, output_len, tool_choice
+            )
+        self.message_history.extend(ChatDocument.to_LLMMessage(response))
         response.metadata.msg_idx = len(self.message_history) - 1
         response.metadata.agent_id = self.id
         # Preserve trail of tool_ids for OpenAI Assistant fn-calls
@@ -622,8 +661,14 @@ class ChatAgent(Agent):
             ):
                 # either the message is a str, or it is a fresh ChatDocument
                 # different from the last message in the history
-                llm_msg = ChatDocument.to_LLMMessage(message)
-                self.message_history.append(llm_msg)
+                llm_msgs = ChatDocument.to_LLMMessage(message, self.oai_tool_calls)
+
+                # process tools if any
+                done_tools = [m.tool_call_id for m in llm_msgs if m.role == Role.TOOL]
+                self.oai_tool_calls = [
+                    t for t in self.oai_tool_calls if t.id not in done_tools
+                ]
+                self.message_history.extend(llm_msgs)
 
         hist = self.message_history
         output_len = self.config.llm.max_output_tokens
@@ -707,18 +752,47 @@ class ChatAgent(Agent):
 
     def _function_args(
         self,
-    ) -> Tuple[Optional[List[LLMFunctionSpec]], str | Dict[str, str]]:
+    ) -> Tuple[
+        Optional[List[LLMFunctionSpec]],
+        str | Dict[str, str],
+        Optional[List[OpenAIToolSpec]],
+        Optional[Dict[str, Dict[str, str] | str]],
+    ]:
+        """Get function/tool spec arguments for OpenAI-compatible LLM API call"""
         functions: Optional[List[LLMFunctionSpec]] = None
         fun_call: str | Dict[str, str] = "none"
+        tools: Optional[List[OpenAIToolSpec]] = None
+        force_tool: Optional[Dict[str, Dict[str, str] | str]] = None
         if self.config.use_functions_api and len(self.llm_functions_usable) > 0:
-            functions = [self.llm_functions_map[f] for f in self.llm_functions_usable]
-            fun_call = (
-                "auto" if self.llm_function_force is None else self.llm_function_force
-            )
-        return functions, fun_call
+            if not self.config.use_tools_api:
+                functions = [
+                    self.llm_functions_map[f] for f in self.llm_functions_usable
+                ]
+                fun_call = (
+                    "auto"
+                    if self.llm_function_force is None
+                    else self.llm_function_force
+                )
+            else:
+                tools = [
+                    OpenAIToolSpec(type="function", function=self.llm_functions_map[f])
+                    for f in self.llm_functions_usable
+                ]
+                force_tool = (
+                    None
+                    if self.llm_function_force is None
+                    else {
+                        "type": "function",
+                        "function": {"name": self.llm_function_force["name"]},
+                    }
+                )
+        return functions, fun_call, tools, force_tool
 
     def llm_response_messages(
-        self, messages: List[LLMMessage], output_len: Optional[int] = None
+        self,
+        messages: List[LLMMessage],
+        output_len: Optional[int] = None,
+        tool_choice: ToolChoiceTypes | Dict[str, str | Dict[str, str]] = "auto",
     ) -> ChatDocument:
         """
         Respond to a series of messages, e.g. with OpenAI ChatCompletion
@@ -748,11 +822,13 @@ class ChatAgent(Agent):
                 stack.enter_context(cm)
             if self.llm.get_stream() and not settings.quiet:
                 console.print(f"[green]{self.indent}", end="")
-            functions, fun_call = self._function_args()
+            functions, fun_call, tools, force_tool = self._function_args()
             assert self.llm is not None
             response = self.llm.chat(
                 messages,
                 output_len,
+                tools=tools,
+                tool_choice=force_tool or tool_choice,
                 functions=functions,
                 function_call=fun_call,
             )
@@ -775,23 +851,24 @@ class ChatAgent(Agent):
             print_response_stats=self.config.show_stats and not settings.quiet,
         )
         chat_doc = ChatDocument.from_LLMResponse(response, displayed=True)
+        self.oai_tool_calls = response.oai_tool_calls or []
+        self.oai_tool_id2call.update(
+            {t.id: t for t in self.oai_tool_calls if t.id is not None}
+        )
         return chat_doc
 
     async def llm_response_messages_async(
-        self, messages: List[LLMMessage], output_len: Optional[int] = None
+        self,
+        messages: List[LLMMessage],
+        output_len: Optional[int] = None,
+        tool_choice: ToolChoiceTypes | Dict[str, str | Dict[str, str]] = "auto",
     ) -> ChatDocument:
         """
         Async version of `llm_response_messages`. See there for details.
         """
         assert self.config.llm is not None and self.llm is not None
         output_len = output_len or self.config.llm.max_output_tokens
-        functions: Optional[List[LLMFunctionSpec]] = None
-        fun_call: str | Dict[str, str] = "none"
-        if self.config.use_functions_api and len(self.llm_functions_usable) > 0:
-            functions = [self.llm_functions_map[f] for f in self.llm_functions_usable]
-            fun_call = (
-                "auto" if self.llm_function_force is None else self.llm_function_force
-            )
+        functions, fun_call, tools, force_tool = self._function_args()
         assert self.llm is not None
 
         streamer = noop_fn
@@ -802,6 +879,8 @@ class ChatAgent(Agent):
         response = await self.llm.achat(
             messages,
             output_len,
+            tools=tools,
+            tool_choice=force_tool or tool_choice,
             functions=functions,
             function_call=fun_call,
         )
@@ -824,6 +903,10 @@ class ChatAgent(Agent):
             print_response_stats=self.config.show_stats and not settings.quiet,
         )
         chat_doc = ChatDocument.from_LLMResponse(response, displayed=True)
+        self.oai_tool_calls = response.oai_tool_calls or []
+        self.oai_tool_id2call.update(
+            {t.id: t for t in self.oai_tool_calls if t.id is not None}
+        )
         return chat_doc
 
     def _render_llm_response(
@@ -847,6 +930,7 @@ class ChatAgent(Agent):
                     if isinstance(response, ChatDocument)
                     else ChatDocument.from_LLMResponse(response, displayed=True)
                 )
+                # TODO: prepend TOOL: or OAI-TOOL: if it's a tool-call
                 print(cached + "[green]" + escape(str(response)))
                 self.callbacks.show_llm_response(
                     content=str(response),
@@ -923,8 +1007,14 @@ class ChatAgent(Agent):
         # If there is a response, then we will have two additional
         # messages in the message history, i.e. the user message and the
         # assistant response. We want to (carefully) remove these two messages.
-        self.message_history.pop() if len(self.message_history) > n_msgs else None
-        self.message_history.pop() if len(self.message_history) > n_msgs else None
+        if len(self.message_history) > n_msgs:
+            msg = self.message_history.pop()
+            self._drop_msg_update_tool_calls(msg)
+
+        if len(self.message_history) > n_msgs:
+            msg = self.message_history.pop()
+            self._drop_msg_update_tool_calls(msg)
+
         return response
 
     async def llm_response_forget_async(self, message: str) -> ChatDocument:
@@ -941,8 +1031,13 @@ class ChatAgent(Agent):
         # If there is a response, then we will have two additional
         # messages in the message history, i.e. the user message and the
         # assistant response. We want to (carefully) remove these two messages.
-        self.message_history.pop() if len(self.message_history) > n_msgs else None
-        self.message_history.pop() if len(self.message_history) > n_msgs else None
+        if len(self.message_history) > n_msgs:
+            msg = self.message_history.pop()
+            self._drop_msg_update_tool_calls(msg)
+
+        if len(self.message_history) > n_msgs:
+            msg = self.message_history.pop()
+            self._drop_msg_update_tool_calls(msg)
         return response
 
     def chat_num_tokens(self, messages: Optional[List[LLMMessage]] = None) -> int: