app.py

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
import os
import logging

def dropdown_update_choices(choices):
    return gr.update(choices=choices, value=None)

# Configure logging
logging.basicConfig(level=logging.INFO) # not working (prints warnings tho)

# os.environ['HF_TOKEN'] = "..."
# hf_token = os.environ.get("HF_TOKEN")

logging.warning('start')

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 process_logs(agent):
    logs = ""
    if hasattr(agent, 'logs'):
        for entry in agent.logs:
            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."

def get_tools():
    return [{"name": tool.name, "description": tool.description} for tool in agent.tools.values()]

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 "Function not found."

def get_tool_description(selected_tool_name, 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
    updated_functions = get_functions()

    tool_names = [tool["name"] for tool in updated_tools]
    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)
    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()

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            with gr.Tab("Chat"):
                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"):

                local_storage = gr.BrowserState(["", ""])
                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]



        with gr.Column():
            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],
                value=tools[0]["name"] if tools else None,
                type="value",
                allow_custom_value=False,
                scale=3
            )

            description_textbox = gr.Textbox(
                label="Tool Description",
                value=get_tool_description(tool_dropdown.value, tools),
                interactive=False,
            )

            slug = gr.Textbox(label="collection slug",value="Mightypeacock/agent-tools-6777c9699c231b7a1e87fa31")
            greet_btn = gr.Button("Load")
            


            gr.Markdown("<center><h2>Functions</h2></center>")
            functions = get_functions()
            function_dropdown = gr.Dropdown(
                label="Select Function",
                choices=list(functions.keys()),
                value=None if not functions.keys() else list(functions.keys())[0],
                type="value",
            )

            code = gr.Code(label="Function Code", language="python")

            tool_dropdown.change(
                fn=get_tool_description,
                inputs=[tool_dropdown, gr.State(tools)],
                outputs=description_textbox,
            )

            function_dropdown.change(
                fn=get_function_code,
                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

    msg.submit(
        respond,
        inputs=[msg, outputbox, chatbot],
        outputs=[msg, outputbox, chatbot, tool_dropdown, function_dropdown, description_textbox, code]
    )

if __name__ == "__main__":
    demo.launch(show_error=True, debug=True)