app.py

import streamlit as st
import os
import base64
import io
from PIL import Image
from pydub import AudioSegment
import IPython
import soundfile as sf
import requests
import pandas as pd  # If you're working with DataFrames
import matplotlib.figure  # If you're using matplotlib figures
import numpy as np

from custom_agent import CustomHfAgent 
from tool_loader import ToolLoader
from tool_config import tool_names   
from app_description import show_app_description


# For Altair charts
import altair as alt

# For Bokeh charts
from bokeh.models import Plot

# For Plotly charts
import plotly.express as px

# For Pydeck charts
import pydeck as pdk


import logging
import streamlit as st
from transformers import load_tool, Agent
from tool_loader import ToolLoader

# Configure the logging settings for transformers
transformers_logger = logging.getLogger("transformers.file_utils")
transformers_logger.setLevel(logging.INFO)  # Set the desired logging level
  

import time  
import torch


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'],
        additional_tools=selected_tools,
        input_params={"max_new_tokens": 192},
    )

    response = agent.run(user_message)

    log_response("Agent Response\n {}".format(response))

    return response

# Declare global variable
global log_enabled
log_enabled = False

def log_response(response):
    if log_enabled:
        with st.chat_message("ai"):
            st.markdown("Agent Response\n {}".format(response))
    print(response)

    


# Create tool loader instance
tool_loader = ToolLoader(tool_names)

st.title("Hugging Face Agent and tools")

## LB https://huggingface.co/spaces/qiantong-xu/toolbench-leaderboard

st.markdown("Welcome to the Hugging Face Agent and Tools app! This app allows you to interact with various tools using the Hugging Face API.")
    
# Create a page with tabs
tabs = st.tabs(["Chat", "URL, Tools and logging", "User Description", "Developers"])

# Tab 1: Chat
with tabs[0]:
    
    # Code for URL and Tools checkboxes 
    #chat_description()
    # Examples for the user perspective
    st.markdown("Stat to chat. e.g. Generate an image of a boat. This will make the agent use the tool text2image to generate an image.")
 
        
# Tab 2: URL and Tools
with tabs[1]:
    #
    app_config()

# Tab 3: User Description
with tabs[2]:
    #
    app_user_description()
    

    
# Tab 4: Developers
with tabs[3]:
         
    # Developer-related content
    st.markdown(''' 

        # Hugging Face Agent and Tools Code Overview
        
        ## Overview
        The provided Python code implements an interactive Streamlit web application that allows users to interact with various tools through the Hugging Face API. The app integrates Hugging Face models and tools, enabling users to perform tasks such as text generation, sentiment analysis, and more.
        
        ## Imports
        The code imports several external libraries and modules, including:
        - `streamlit`: For building the web application.
        - `os`: For interacting with the operating system.
        - `base64`, `io`, `Image` (from `PIL`), `AudioSegment` (from `pydub`), `IPython`, `sf`: For handling images and audio.
        - `requests`: For making HTTP requests.
        - `pandas`: For working with DataFrames.
        - `matplotlib.figure`, `numpy`: For visualization.
        - `altair`, `Plot` (from `bokeh.models`), `px` (from `plotly.express`), `pdk` (from `pydeck`): For different charting libraries.
        - `time`: For handling time-related operations.
        - `transformers`: For loading tools and agents.
        
        ## ToolLoader Class
        The `ToolLoader` class is responsible for loading tools based on their names. It has methods to load tools from a list of tool names and handles potential errors during loading.
        
        ## CustomHfAgent Class
        The `CustomHfAgent` class extends the base `Agent` class from the `transformers` module. It is designed to interact with a remote inference API and includes methods for generating text based on a given prompt.
        
        ## Tool Loading and Customization
        - Tool names are defined in the `tool_names` list.
        - The `ToolLoader` instance (`tool_loader`) loads tools based on the provided names.
        - The `CustomHfAgent` instance (`agent`) is created with a specified URL endpoint, token, and additional tools.
        - New tools can be added by appending their names to the `tool_names` list.
        
        ## Streamlit App
        The Streamlit app is structured as follows:
        1. Tool selection dropdown for choosing the inference URL.
        2. An expander for displaying tool descriptions.
        3. An expander for selecting tools.
        4. Examples and instructions for the user.
        5. A chat interface for user interactions.
        6. Handling of user inputs, tool selection, and agent responses.
        
        ## Handling of Responses
        The code handles various types of responses from the agent, including images, audio, text, DataFrames, and charts. The responses are displayed in the Streamlit app based on their types.
        
        ## How to Run
        1. Install required dependencies with `pip install -r requirements.txt`.
        2. Run the app with `streamlit run <filename.py>`.
        
        ## Notes
        - The code emphasizes customization and extensibility, allowing developers to easily add new tools and interact with the Hugging Face API.
        - Ensure proper configuration, such as setting the Hugging Face token as an environment variable.

      ''')


        # Display logs in the frontend
    logs_expander = st.expander("Logs")
    with logs_expander:
        log_output = st.empty()
    
        # Custom logging handler to append log messages to the chat
        class ChatHandler(logging.Handler):
            def __init__(self):
                super().__init__()
    
            def emit(self, record):
                log_message = self.format(record)
                with st.chat_message("ai"):
                    st.markdown(f"Log: {log_message}")
    
        # Add the custom handler to the transformers_logger
        chat_handler = ChatHandler()
        transformers_logger.addHandler(chat_handler)
    
        # Function to update logs in the frontend
        def update_logs():
            log_output.code("")  # Clear previous logs
            # Do nothing here since logs are appended to the chat
    
        # Update logs when the button is clicked
        if st.button("Update Logs"):
            update_logs()
    
    
# Chat code (user input, agent responses, etc.)
if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

with st.chat_message("assistant"):
    st.markdown("Hello there! How can I assist you today?")

if user_message := st.chat_input("Enter message"):
    st.chat_message("user").markdown(user_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]
    # 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:
            st.warning("The agent's response is None. Please try again. Generate an image of a flying horse.")
        elif isinstance(response, Image.Image):
            st.image(response)
        elif isinstance(response, AudioSegment):
            st.audio(response)
        elif isinstance(response, int):
            st.markdown(response)
        elif isinstance(response, str):
            if "emojified_text" in response:
                st.markdown(f"{response['emojified_text']}")
            else:
                st.markdown(response)
        elif isinstance(response, list):
            for item in response:
                st.markdown(item)  # Assuming the list contains strings
        elif isinstance(response, pd.DataFrame):
            st.dataframe(response)
        elif isinstance(response, pd.Series):
            st.table(response.iloc[0:10])
        elif isinstance(response, dict):
            st.json(response)
        elif isinstance(response, st.graphics_altair.AltairChart):
            st.altair_chart(response)
        elif isinstance(response, st.graphics_bokeh.BokehChart):
            st.bokeh_chart(response)
        elif isinstance(response, st.graphics_graphviz.GraphvizChart):
            st.graphviz_chart(response)
        elif isinstance(response, st.graphics_plotly.PlotlyChart):
            st.plotly_chart(response)
        elif isinstance(response, st.graphics_pydeck.PydeckChart):
            st.pydeck_chart(response)
        elif isinstance(response, matplotlib.figure.Figure):
            st.pyplot(response)
        elif isinstance(response, streamlit.graphics_vega_lite.VegaLiteChart):
            st.vega_lite_chart(response)
        else:
            st.warning("Unrecognized response type. Please try again. e.g. Generate an image of a flying horse.")

    st.session_state.messages.append({"role": "assistant", "content": response})