Update app.py

app.py

@@ -28,7 +28,6 @@ import time
 from transformers import load_tool, Agent
 import torch
 
-
 class ToolLoader:
     def __init__(self, tool_names):
         self.tools = self.load_tools(tool_names)
@@ -41,8 +40,6 @@ class ToolLoader:
                 loaded_tools.append(tool)
             except Exception as e:
                 print(f"Error loading tool '{tool_name}': {e}")
-                # Handle the error as needed, e.g., continue with other tools or take corrective action
-
         return loaded_tools
 
 class CustomHfAgent(Agent):
@@ -79,8 +76,19 @@ class CustomHfAgent(Agent):
                 return result[: -len(stop_seq)]
         return result
 
-def load_tools(tool_names):
-    return [load_tool(tool_name) for tool_name in tool_names]
+def handle_submission(user_message, selected_tools, url_endpoint):
+    agent = CustomHfAgent(
+        url_endpoint=url_endpoint,
+        token=os.environ['HF_token'],
+        additional_tools=selected_tools,
+        input_params={"max_new_tokens": 192},
+    )
+
+    response = agent.run(user_message)
+
+    print("Agent Response\n {}".format(response))
+
+    return response
 
 # Define the tool names to load
 tool_names = [
@@ -88,44 +96,37 @@ tool_names = [
     "Chris4K/text-generation-tool",
     "Chris4K/sentiment-tool", 
     "Chris4K/token-counter-tool", 
-	"Chris4K/most-downloaded-model",
-	"Chris4K/rag-tool",
-	"Chris4K/word-counter-tool",
-	"Chris4K/sentence-counter-tool",
-	"Chris4K/EmojifyTextTool",
-	"Chris4K/NamedEntityRecognitionTool",
-	"Chris4K/TextDownloadTool",
-	"Chris4K/source-code-retriever-tool",
-	"Chris4K/text-to-image",
-	"Chris4K/text-to-video",
-	"Chris4K/image-transformation",
-	"Chris4K/latent-upscaler-tool"
+    "Chris4K/most-downloaded-model",
+    "Chris4K/rag-tool",
+    "Chris4K/word-counter-tool",
+    "Chris4K/sentence-counter-tool",
+    "Chris4K/EmojifyTextTool",
+    "Chris4K/NamedEntityRecognitionTool",
+    "Chris4K/TextDownloadTool",
+    "Chris4K/source-code-retriever-tool",
+    "Chris4K/text-to-image",
+    "Chris4K/text-to-video",
+    "Chris4K/image-transformation",
+    "Chris4K/latent-upscaler-tool"
     # More cool tools to come
 ]
 
 # Create tool loader instance
 tool_loader = ToolLoader(tool_names)
 
-# Define the callback function to handle the form submission
-def handle_submission(user_message, selected_tools):
-    agent = CustomHfAgent(
-        url_endpoint="https://api-inference.huggingface.co/models/bigcode/starcoder",
-        token=os.environ['HF_token'],
-        additional_tools=selected_tools,
-        input_params={"max_new_tokens": 192},
-    )
-
-    response = agent.run(user_message)
-
-    print("Agent Response\n {}".format(response))
-
-    return response
-
 st.title("Hugging Face Agent and tools")
 
+# Add a dropdown for selecting the inference URL
+url_endpoint = st.selectbox("Select Inference URL", [
+    "https://api-inference.huggingface.co/models/bigcode/starcoder",
+    "https://api-inference.huggingface.co/models/bigcode/starcode",
+    "https://api-inference.huggingface.co/models/gpt"
+])
+ 
 tool_checkboxes = [st.checkbox(f"{tool.name} --- {tool.description} ") for tool in tool_loader.tools]
 
 
+
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
@@ -142,7 +143,9 @@ if user_message := st.chat_input("Enter message"):
     st.session_state.messages.append({"role": "user", "content": user_message})
 
     selected_tools = [tool_loader.tools[idx] for idx, checkbox in enumerate(tool_checkboxes) if checkbox]
-    response = handle_submission(user_message, selected_tools)
+    # Handle submission with the selected inference URL
+    response = handle_submission(user_message, selected_tools, url_endpoint)
+
 
     with st.chat_message("assistant"):
         if response is None: