fix: optimized clearing registry and tool calls pool

src/writer/ai.py

@@ -896,6 +896,9 @@ def _clear_callable_registry(self):
         self._callable_registry = {}
 
     def _clear_ongoing_tool_calls(self):
+        """
+        Clear ongoing tool calls after they've been processed
+        """
         self._ongoing_tool_calls = {}
 
     def _clear_tool_calls_helpers(self):
@@ -1328,11 +1331,10 @@ def _process_response_data(
                     )
                     logger.debug(f"Received response – {follow_up_response}")
 
-                    # Clear buffer and callable registry for the completed tool call
-                    self._clear_tool_calls_helpers()
-
                     # Call the function recursively to either process a new tool call
                     # or return the message if no tool calls are requested
+
+                    self._clear_ongoing_tool_calls()
                     return self._process_response_data(
                         follow_up_response,
                         passed_messages=passed_messages,
@@ -1382,9 +1384,8 @@ def _process_stream_response(
                         )
                     )
 
-                    # Clear buffer and callable registry for the completed tool call
                     try:
-                        self._clear_tool_calls_helpers()
+                        self._clear_ongoing_tool_calls()
                         yield from self._process_stream_response(
                             response=follow_up_response,
                             passed_messages=passed_messages,
@@ -1449,14 +1450,19 @@ def complete(
                 )
             )
 
-        return self._process_response_data(
+        response = self._process_response_data(
             response_data,
             passed_messages=passed_messages,
             request_model=request_model,
             request_data=request_data,
             max_depth=max_tool_depth
             )
 
+        # Clear buffer and callable registry for the completed tool call
+        self._clear_tool_calls_helpers()
+
+        return response
+
     def stream_complete(
             self,
             config: Optional['ChatOptions'] = None,
@@ -1512,6 +1518,8 @@ def stream_complete(
             max_depth=max_tool_depth
         )
 
+        # Clear buffer and callable registry for the completed tool call
+        self._clear_tool_calls_helpers()
         response.close()
 
     @property