Update app.py

app.py

@@ -39,7 +39,7 @@ class ToolLoader:
                 tool = load_tool(tool_name)
                 loaded_tools.append(tool)
             except Exception as e:
-                print(f"Error loading tool '{tool_name}': {e}")
+                log_response(f"Error loading tool '{tool_name}': {e}")
         return loaded_tools
 
 class CustomHfAgent(Agent):
@@ -64,12 +64,12 @@ class CustomHfAgent(Agent):
         response = requests.post(self.url_endpoint, json=inputs, headers=headers)
 
         if response.status_code == 429:
-            print("Getting rate-limited, waiting a tiny bit before trying again.")
+            log_response("Getting rate-limited, waiting a tiny bit before trying again.")
             time.sleep(1)
             return self._generate_one(prompt)
         elif response.status_code != 200:
             raise ValueError(f"Errors {inputs} {response.status_code}: {response.json()}")
-        print(response)
+        log_response(response)
         result = response.json()[0]["generated_text"]
         for stop_seq in stop:
             if result.endswith(stop_seq):
@@ -77,6 +77,11 @@ class CustomHfAgent(Agent):
         return result
 
 def handle_submission(user_message, selected_tools, url_endpoint):
+    
+    log_response("User input \n {}".format(user_message))
+    log_response("selected_tools \n {}".format(selected_tools))
+    log_response("url_endpoint \n {}".format(url_endpoint))
+    
     agent = CustomHfAgent(
         url_endpoint=url_endpoint,
         token=os.environ['HF_token'],
@@ -86,10 +91,18 @@ def handle_submission(user_message, selected_tools, url_endpoint):
 
     response = agent.run(user_message)
 
-    print("Agent Response\n {}".format(response))
+    log_response("Agent Response\n {}".format(response))
 
     return response
 
+    
+def log_response(response):
+    if log_enabled:
+        with st.chat_message("ai"):
+            st.markdown("Agent Response\n {}".format(response))
+    print(response)
+
+    
 # Define the tool names to load
 tool_names = [
     "Chris4K/random-character-tool",
@@ -167,6 +180,8 @@ elif tabs[1]:
         "https://api-inference.huggingface.co/models/gpt2"
     ])
 
+        # Add a checkbox for enabling logging
+    log_enabled = st.checkbox("Enable Logging")
     
     tool_checkboxes = [st.checkbox(f"{tool.name} --- {tool.description} ") for tool in tool_loader.tools]