refactor

app.py

@@ -1,34 +1,90 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-import random
-import time
-from io import StringIO
-import sys
-import pprint
-import inspect
-from HFHub import login_for_publication, litellm_api_keys, load_collection_from_space
-from HFHub import createAgent
-from typing import Dict
+from typing import Dict, List, Set
 import os
 import logging
-
-def dropdown_update_choices(choices):
-    return gr.update(choices=choices, value=None)
-
+import inspect
+from huggingface_hub import login
+from smolagents import ToolCollection, CodeAgent, Tool, HfApiModel
 # Configure logging
-logging.basicConfig(level=logging.INFO) # not working (prints warnings tho)
+logging.basicConfig(level=logging.INFO)
+logging.warning('Starting application...')
+
+# Global variables
+publishtoken = None
+tool_Collections: dict[str, ToolCollection] = {}
+loaded_tools: Set[Tool] = set()
+
+# API keys for litellm providers
+litellm_api_keys = {
+    'xai': '',
+    'HF': '',
+    'grok': '',
+    'anthropic': '',
+    'openAI': '',
+}
+
+def all_tools_names() -> List[str]:
+    """Return a list of all tool names from loaded collections."""
+    all_tools = []
+    for collection in tool_Collections.values():
+        if isinstance(collection, ToolCollection):
+            all_tools.extend(tool.name for tool in collection.tools)
+        else:
+            all_tools.extend(collection)
+    return all_tools
+
+def filter_tools(tools):
+    """Filter out base tools from a list of tools."""
+    if tools is None:
+        logging.warning("Received None for tools, defaulting to an empty list.")
+        tools = []
+
+    base_tools_names = ['web search']
+    return [tool for tool in tools if tool not in base_tools_names]
+
+def load_collection_from_space(agent: CodeAgent, collection_slug: str) -> List[str]:
+    """Load a collection of tools from a Hugging Face space."""
+    if collection_slug not in tool_Collections:
+        tool_collection = ToolCollection(
+            collection_slug=collection_slug,
+            trust_remote_code=True
+        )
+
+        tool_Collections[collection_slug] = [tool.name for tool in tool_collection.tools]
+
+        for tool in tool_collection.tools:
+            if agent.tools.get(tool.name) is None:
+                agent.tools[tool.name] = tool
+                loaded_tools.add(tool)
+            else:
+                agent.tools[tool.name] = tool
+
+    return all_tools_names()
+
+def createAgent() -> CodeAgent:
+    """Create and return a CodeAgent instance."""
+    agent = CodeAgent(
+        tools=filter_tools(list(loaded_tools)),
+        model=HfApiModel(),
+        additional_authorized_imports=["smolagents", "subprocess", "typing", "os", "inspect", "open", "requests"],
+        add_base_tools=True,
+        planning_interval=None,
+    )
 
-# os.environ['HF_TOKEN'] = "..."
-# hf_token = os.environ.get("HF_TOKEN")
+    # Add base tools to the collection
+    for tool in agent.tools:
+        if tool not in loaded_tools:
+            if "base tools" not in tool_Collections:
+                tool_Collections["base tools"] = []
+            tool_Collections["base tools"].append(tool)
 
-logging.warning('start')
+    return agent
 
+# Initialize the agent
 agent = createAgent()
 
-def update_agent(collection_slug: str):
-    load_collection_from_space(agent, collection_slug=collection_slug)
-    return refresh_ui_elements()  # Update the UI elements after loading the collection
-
+def dropdown_update_choices(choices):
+    return gr.update(choices=choices, value=None)
 
 def process_logs(agent):
     logs = ""
@@ -37,8 +93,7 @@ def process_logs(agent):
             if hasattr(entry, 'llm_output'):
                 logs += str(entry.llm_output) + "\n"
         return logs
-    else:
-        return "The agent object does not have a valid 'logs' attribute or the attribute is not a list."
+    return "No logs available."
 
 def get_tools():
     return [{"name": tool.name, "description": tool.description} for tool in agent.tools.values()]
@@ -46,113 +101,94 @@ def get_tools():
 def get_functions():
     return agent.python_executor.custom_tools
 
-
-
 def get_function_code(selected_function_name):
     func = get_functions().get(selected_function_name)
     if func:
         try:
             return inspect.getsource(func)
-
         except OSError:
-            return "Source code not available for this function."
+            return "Source code not available."
     return "Function not found."
 
-def get_tool_description(selected_tool_name, tools):
+def get_tool_description(selected_tool_name):
+    
+    tools = get_tools()
+    print("Selected tool name:", selected_tool_name)
+    print("Tools:",tools )
     for tool in tools:
         if tool["name"] == selected_tool_name:
             return tool["description"]
     return "No description available."
 
-def get_tools_name():
-    return [tool["name"] for tool in get_tools()]
-    
-    
 def refresh_ui_elements():
-    updated_tools = get_tools()  # Get the updated tools after loading the collection
+    print("Refreshing UI elements...")
+    updated_tools = get_tools()
     updated_functions = get_functions()
 
     tool_names = [tool["name"] for tool in updated_tools]
+    print("Tool names:", tool_names)
     function_names = list(updated_functions.keys())
 
-    print(f"function_names: {function_names}")
-
     current_tool = tool_names[0] if tool_names else None
     current_function = function_names[0] if function_names else None
 
-    tool_description = get_tool_description(current_tool, updated_tools)
+    tool_description = get_tool_description(current_tool)
     function_code = get_function_code(current_function) if current_function else ""
 
-    # Update dropdown choices
     tool_dropdown_update = dropdown_update_choices(tool_names)
     function_dropdown_update = dropdown_update_choices(function_names)
-    return (
-        tool_dropdown_update,
-        function_dropdown_update,
-        tool_description,
-        function_code
-    )
 
-def load_collection_andUploadUI(collection_slug: str ):
-    # load_collection_from_space(agent,collection_slug)
-    refresh_ui_elements()
+    return tool_dropdown_update, function_dropdown_update, tool_description, function_code
+
+def update_agent(collection_slug: str):
+    load_collection_from_space(agent, collection_slug=collection_slug)
+    return refresh_ui_elements()
+
+def respond(message, console_output, chat_history):
+    try:
+        bot_message = agent.run(message)
+        new_console_output = process_logs(agent)
+
+        chat_history.extend([
+            {"role": "user", "content": message},
+            {"role": "assistant", "content": bot_message}
+        ])
+
+        updated_console = console_output + f"\nQuery: {message}\nLogs: {new_console_output}"
+        ui_updates = refresh_ui_elements()
+
+        return "", updated_console, chat_history, *ui_updates
+    except Exception as e:
+        logging.error(f"Error in respond function: {e}")
+        return f"An error occurred: {str(e)}", console_output, chat_history, None, None, None, None
 
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             with gr.Tab("Chat"):
-                gr.Markdown("<center><h1>smolAgent chat</h1></center>")
+                gr.Markdown("<center><h1>smolAgent Chat</h1></center>")
                 chatbot = gr.Chatbot(type="messages")
                 msg = gr.Textbox()
-
-                # send_button = gr.Button("Send")  # Send button
-                # send_button.click(respond, [msg, chatbot], [msg, chatbot])
-
                 clear = gr.ClearButton([msg, chatbot])
 
             with gr.Tab("Console"):
-                outputbox = gr.Textbox(
-                    lines=25,
-                    scale=1,
-                    interactive=False
-                )
-
-            with gr.Tab("config"):
+                outputbox = gr.Textbox(lines=25, scale=1, interactive=False)
 
-                local_storage = gr.BrowserState(["", ""])
+            with gr.Tab("Config"):
                 gr.Markdown("## Configure litellm API Keys")
-                api_key_inputs = {}
-                for provider in litellm_api_keys.keys():
-                    with gr.Row():
-                        api_key_inputs[provider] = gr.Textbox(
-                            label=f'{provider} API Key',
-                            placeholder='Enter key',
-                            type='password',
-                            value=litellm_api_keys[provider]
-                        )
-
-#TODO : save keys to localstorage
-
-                        # api_key_inputs[provider]
-                # save_keys_button = gr.Button('Save Configuration')
-
-
-                #         @gr.on(api_key_inputs[provider].change, inputs=[api_key_inputs[provider]], outputs=[local_storage])
-                #         def save_to_local_storage(text):
-                #             return [text]
-
-                # # Load from local storage when page loads
-                # @demo.load(inputs=[local_storage], outputs=[input_text])
-                # def load_from_local_storage(saved_values):
-                #     return saved_values[0]
-
-
+                api_key_inputs = {
+                    provider: gr.Textbox(
+                        label=f'{provider} API Key',
+                        placeholder='Enter key',
+                        type='password',
+                        value=litellm_api_keys[provider]
+                    ) for provider in litellm_api_keys
+                }
 
         with gr.Column():
-            gr.Markdown("<center><h1>tool collection</h1></center>")
-
-
+            gr.Markdown("<center><h1>Tool Collection</h1></center>")
             tools = get_tools()
+
             tool_dropdown = gr.Dropdown(
                 show_label=False,
                 choices=[tool["name"] for tool in tools],
@@ -164,13 +200,12 @@ with gr.Blocks() as demo:
 
             description_textbox = gr.Textbox(
                 label="Tool Description",
-                value=get_tool_description(tool_dropdown.value, tools),
+                value=get_tool_description(tool_dropdown.value),
                 interactive=False,
             )
 
-            slug = gr.Textbox(label="collection slug",value="Mightypeacock/agent-tools-6777c9699c231b7a1e87fa31")
+            slug = gr.Textbox(label="Collection Slug", value="Mightypeacock/agent-tools-6777c9699c231b7a1e87fa31")
             greet_btn = gr.Button("Load")
-            
 
 
             gr.Markdown("<center><h2>Functions</h2></center>")
@@ -186,7 +221,7 @@ with gr.Blocks() as demo:
 
             tool_dropdown.change(
                 fn=get_tool_description,
-                inputs=[tool_dropdown, gr.State(tools)],
+                inputs=[tool_dropdown],
                 outputs=description_textbox,
             )
 
@@ -195,40 +230,12 @@ with gr.Blocks() as demo:
                 inputs=function_dropdown,
                 outputs=code,
             )
-
-            greet_btn.click(fn=update_agent, inputs=slug,
-                      
-                         outputs=[tool_dropdown, function_dropdown, description_textbox, code],
-    #   outputs=tool_dropdown,
-                            # outputs=None,
-                            api_name="load_HF_Collection")
-
-    def respond(message, console_output, chat_history):
-        try:
-            print(f"Received message: {message}")
-
-            if not isinstance(message, str):
-                message = str(message)
-
-            bot_message = agent.run(message)
-            new_console_output = process_logs(agent)
-
-            print(f"Agent response: {bot_message}")
-            print(f"Console output: {new_console_output}")
-
-            chat_history.extend([
-                {"role": "user", "content": message},
-                {"role": "assistant", "content": bot_message}
-            ])
-
-            updated_console = console_output + f"\nQuery: {message}\nLogs: {new_console_output}"
-
-            tool_dropdown_update, function_dropdown_update, description_update, code_update = refresh_ui_elements()
-            return "", updated_console, chat_history, tool_dropdown_update, function_dropdown_update, description_update, code_update
-
-        except Exception as e:
-            print(f"Error in respond function: {e}")
-            return f"An error occurred: {str(e)}", console_output, chat_history, None, None, None, None
+            greet_btn.click(
+                fn=update_agent,
+                inputs=slug,
+                outputs=[tool_dropdown, function_dropdown, description_textbox, code],
+                api_name="load_HF_Collection"
+            )
 
     msg.submit(
         respond,