From 1582017ad4e0225416fd27392969cbeeb872c5d3 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Sun, 29 Dec 2024 17:00:15 +0000
Subject: [PATCH] fix: implement _remember_format in ToolUsage class (#1815)

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/agents/crew_agent_executor.py |  8 +++-----
 src/crewai/tools/tool_usage.py           | 24 +++++++++++++++++-------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/src/crewai/agents/crew_agent_executor.py b/src/crewai/agents/crew_agent_executor.py
index 23a4d3a5d..c4c5477b9 100644
--- a/src/crewai/agents/crew_agent_executor.py
+++ b/src/crewai/agents/crew_agent_executor.py
@@ -127,10 +127,6 @@ class CrewAgentExecutor(CrewAgentExecutorMixin):
                         "Invalid response from LLM call - None or empty."
                     )
 
-                # Remember output format before parsing
-                if hasattr(self.task, 'output_format'):
-                    self._remember_format()
-
                 if not self.use_stop_words:
                     try:
                         self._format_answer(answer)
@@ -192,7 +188,9 @@ class CrewAgentExecutor(CrewAgentExecutorMixin):
                     )
 
         except OutputParserException as e:
-            error_msg = "Error on parsing tool." if "Error parsing tool usage" in str(e) else e.error
+            error_msg = e.error
+            if "Error parsing tool usage" in str(e):
+                error_msg = "Error on parsing tool."
             self.messages.append({"role": "user", "content": error_msg})
             if self.iterations > self.log_error_after:
                 self._printer.print(
diff --git a/src/crewai/tools/tool_usage.py b/src/crewai/tools/tool_usage.py
index 532587ced..f68162e51 100644
--- a/src/crewai/tools/tool_usage.py
+++ b/src/crewai/tools/tool_usage.py
@@ -145,6 +145,10 @@ class ToolUsage:
         from_cache = False
 
         result = None  # type: ignore # Incompatible types in assignment (expression has type "None", variable has type "str")
+        
+        # Remember output format if needed
+        self._remember_format()
+        
         # check if cache is available
         if self.tools_handler.cache:
             result = self.tools_handler.cache.read(  # type: ignore # Incompatible types in assignment (expression has type "str | None", variable has type "str")
@@ -259,14 +263,20 @@ class ToolUsage:
         return result
 
     def _should_remember_format(self) -> bool:
-        return self.task.used_tools % self._remember_format_after_usages == 0
+        """Check if we should remember the output format based on usage count."""
+        return self.task.used_tools >= self._remember_format_after_usages
 
-    def _remember_format(self, result: str) -> None:
-        result = str(result)
-        result += "\n\n" + self._i18n.slice("tools").format(
-            tools=self.tools_description, tool_names=self.tools_names
-        )
-        return result  # type: ignore # No return value expected
+    def _remember_format(self, result: str = None) -> str:
+        """Remember the output format for the task and optionally format the result."""
+        if hasattr(self.task, 'output_format') and self._should_remember_format():
+            self.task.output_format = True
+        
+        if result is not None:
+            result = str(result)
+            result += "\n\n" + self._i18n.slice("tools").format(
+                tools=self.tools_description, tool_names=self.tools_names
+            )
+        return result if result is not None else ""
 
     def _check_tool_repeated_usage(
         self, calling: Union[ToolCalling, InstructorToolCalling]