From e827c2bdc1512443a4146cfe92728a762e8cd178 Mon Sep 17 00:00:00 2001
From: mmikita95 <hello@mmikita.me>
Date: Wed, 16 Oct 2024 16:30:35 +0400
Subject: [PATCH] fix: optimized clearing registry and tool calls pool

---
 src/writer/ai.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/writer/ai.py b/src/writer/ai.py
index f01cd67c9..eba5d9768 100644
--- a/src/writer/ai.py
+++ b/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,7 +1450,7 @@ def complete(
                 )
             )
 
-        return self._process_response_data(
+        response = self._process_response_data(
             response_data,
             passed_messages=passed_messages,
             request_model=request_model,
@@ -1457,6 +1458,11 @@ def complete(
             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