improve number of messages output

agent.py

@@ -274,8 +274,7 @@ class SantaAgent:
             ChatMessage(role="user", content=user_prompt),
         ]
 
-        non_tool_count = 0
-
+        consequtive_non_tool_call_count = 0
         while True:
             response = self.client.chat.completions.create(
                 messages=messages,
@@ -283,15 +282,24 @@ class SantaAgent:
                 tools=self.tools,
                 tool_choice="auto",
             )
+
+            tool_calls = response.choices[0].message.tool_calls
+
+            # Reduce the number of non-toolcall messages
+            if not tool_calls:
+                consequtive_non_tool_call_count += 1
+                if consequtive_non_tool_call_count >= 2:
+                    break
+            else:
+                consequtive_non_tool_call_count = 0
+
             messages.append(response.choices[0].message.to_dict())
             content = response.choices[0].message.content
             if content is not None:
                 gradio_messages.append(ChatMessage(role="assistant", content=content))
-            tool_calls = response.choices[0].message.tool_calls
 
             should_stop = False
             if tool_calls:
-                non_tool_count = 0
                 for tool_call in tool_calls:
                     arguments = json.loads(tool_call.function.arguments)
                     if tool_call.function.name == "buy_item":
@@ -346,12 +354,6 @@ class SantaAgent:
                     if not should_stop:
                         gradio_messages.append(ChatMessage(role="assistant", content=output, metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
 
-            else:
-                non_tool_count += 1
-
-                if non_tool_count >= 2:
-                    break
-
             if should_stop or len(messages) > 10:
                 break
         return messages, gradio_messages

app.py

@@ -48,7 +48,7 @@ def run_testing(user_prompt, invariant_api_key):
 
     agent_params = {"system_prompt": user_prompt}
 
-    yield f'Running Test 0 of {TOTAL_TESTS}. Please wait.', '', 'button-loading'
+    yield f'Running Test 1 of {TOTAL_TESTS}. Please wait.', '', 'button-loading'
 
     env={
         "INVARIANT_API_KEY": invariant_api_key,
@@ -137,7 +137,7 @@ with gr.Blocks(
                 gr.Markdown("""
                     ## Get an API Key
                     * [Create an account here](https://explorer.invariantlabs.ai/settings) by clicking 'Sign In', and then the GitHub icon.
-                    * Click on `Get API Key` to get your Invariant API key.
+                    * Click on `Get API Key` below to get your Invariant API key.
                     * Paste the API key in the text box below.
                     """
                 )