Update app.py

app.py

@@ -1,159 +1,144 @@
-import os
-from typing import Dict, Any
-import asyncio
+from typing import Dict, Any, List, Tuple, Optional
 
-# Check if the OPENAI_API_KEY is set
-if "OPENAI_API_KEY" not in os.environ:
-    raise ValueError("No OPENAI_API_KEY found. Please set the environment variable.")
+from llama_index.llama_pack.base import BaseLlamaPack
+from llama_index.llms import OpenAI
+from llama_index.agent import ReActAgent
 
-# Create a new event loop
-loop = asyncio.new_event_loop()
+from llama_hub.tools.arxiv import ArxivToolSpec
+from llama_hub.tools.wikipedia import WikipediaToolSpec
+import functools
 
-# Set the event loop as the current event loop
-asyncio.set_event_loop(loop)
+from io import StringIO
+import sys
 
-from llama_index import (
-    VectorStoreIndex,
-    ServiceContext,
-    download_loader,
-)
-from llama_index.llama_pack.base import BaseLlamaPack
-from llama_index.llms import OpenAI
+
+SUPPORTED_TOOLS = {
+    "arxiv_search_tool": ArxivToolSpec,
+    "wikipedia": WikipediaToolSpec,
+}
 
 
-class StreamlitChatPack(BaseLlamaPack):
-    """Streamlit chatbot pack."""
+class Capturing(list):
+    """To capture the stdout from ReActAgent.chat with verbose=True. Taken from
+    https://stackoverflow.com/questions/16571150/\
+        how-to-capture-stdout-output-from-a-python-function-call
+    """
+
+    def __enter__(self):
+        self._stdout = sys.stdout
+        sys.stdout = self._stringio = StringIO()
+        return self
+
+    def __exit__(self, *args):
+        self.extend(self._stringio.getvalue().splitlines())
+        del self._stringio  # free up some memory
+        sys.stdout = self._stdout
+
+
+class GradioReActAgentPack(BaseLlamaPack):
+    """Gradio chatbot to chat with a ReActAgent pack."""
 
     def __init__(
         self,
-        wikipedia_page: str = "Snowflake Inc.",
-        run_from_main: bool = False,
+        tools_list: Optional[List[str]] = [k for k in SUPPORTED_TOOLS.keys()],
         **kwargs: Any,
     ) -> None:
         """Init params."""
-        if not run_from_main:
-            raise ValueError(
-                "Please run this llama-pack directly with "
-                "`streamlit run [download_dir]/streamlit_chatbot/base.py`"
-            )
+        try:
+            from ansi2html import Ansi2HTMLConverter
+        except ImportError:
+            raise ImportError("Please install ansi2html via `pip install ansi2html`")
+
+        tools = []
+        for t in tools_list:
+            try:
+                tools.append(SUPPORTED_TOOLS[t]())
+            except KeyError:
+                raise KeyError(f"Tool {t} is not supported.")
+        self.tools = tools
+
+        self.llm = OpenAI(model="gpt-4-1106-preview", max_tokens=2000)
+        self.agent = ReActAgent.from_tools(
+            tools=functools.reduce(
+                lambda x, y: x.to_tool_list() + y.to_tool_list(), self.tools
+            ),
+            llm=self.llm,
+            verbose=True,
+        )
 
-        self.wikipedia_page = wikipedia_page
+        self.thoughts = ""
+        self.conv = Ansi2HTMLConverter()
 
     def get_modules(self) -> Dict[str, Any]:
         """Get modules."""
-        return {}
+        return {"agent": self.agent, "llm": self.llm, "tools": self.tools}
+
+    def _handle_user_message(self, user_message, history):
+        """Handle the user submitted message. Clear message box, and append
+        to the history."""
+        return "", history + [(user_message, "")]
+
+    def _generate_response(
+        self, chat_history: List[Tuple[str, str]]
+    ) -> Tuple[str, List[Tuple[str, str]]]:
+        """Generate the response from agent, and capture the stdout of the
+        ReActAgent's thoughts.
+        """
+        with Capturing() as output:
+            response = self.agent.stream_chat(chat_history[-1][0])
+        ansi = "\n========\n".join(output)
+        html_output = self.conv.convert(ansi)
+        for token in response.response_gen:
+            chat_history[-1][1] += token
+            yield chat_history, str(html_output)
+
+    def _reset_chat(self) -> Tuple[str, str]:
+        """Reset the agent's chat history. And clear all dialogue boxes."""
+        # clear agent history
+        self.agent.reset()
+        return "", "", ""  # clear textboxes
 
     def run(self, *args: Any, **kwargs: Any) -> Any:
         """Run the pipeline."""
-        import streamlit as st
-        from streamlit_pills import pills
-
-        st.set_page_config(
-            page_title=f"Chat with {self.wikipedia_page}'s Wikipedia page, powered by LlamaIndex",
-            page_icon="🦙",
-            layout="centered",
-            initial_sidebar_state="auto",
-            menu_items=None,
-        )
-
-        if "messages" not in st.session_state:  # Initialize the chat messages history
-            st.session_state["messages"] = [
-                {"role": "assistant", "content": "Ask me a question about Snowflake!"}
-            ]
+        import gradio as gr
 
-        st.title(
-            f"Chat with {self.wikipedia_page}'s Wikipedia page, powered by LlamaIndex 💬🦙"
+        demo = gr.Blocks(
+            theme="gstaff/xkcd",
+            css="#box { height: 420px; overflow-y: scroll !important}",
         )
-        st.info(
-            "This example is powered by the **[Llama Hub Wikipedia Loader](https://llamahub.ai/l/wikipedia)**. Use any of [Llama Hub's many loaders](https://llamahub.ai/) to retrieve and chat with your data via a Streamlit app.",
-            icon="ℹ️",
-        )
-
-        def add_to_message_history(role, content):
-            message = {"role": role, "content": str(content)}
-            st.session_state["messages"].append(
-                message
-            )  # Add response to message history
-
-        @st.cache_resource
-        def load_index_data():
-            WikipediaReader = download_loader(
-                "WikipediaReader", custom_path="local_dir"
-            )
-            loader = WikipediaReader()
-            docs = loader.load_data(pages=[self.wikipedia_page])
-            service_context = ServiceContext.from_defaults(
-                llm=OpenAI(model="gpt-3.5-turbo", temperature=0.5)
-            )
-            index = VectorStoreIndex.from_documents(
-                docs, service_context=service_context
+        with demo:
+            gr.Markdown(
+                "# Gradio ReActAgent Powered by LlamaIndex and LlamaHub 🦙\n"
+                "This Gradio app is powered by LlamaIndex's `ReActAgent` with\n"
+                "OpenAI's GPT-4-Turbo as the LLM. The tools are listed below.\n"
+                "## Tools\n"
+                "- [ArxivToolSpec](https://llamahub.ai/l/tools-arxiv)\n"
+                "- [WikipediaToolSpec](https://llamahub.ai/l/tools-wikipedia)"
             )
-            return index
-
-        index = load_index_data()
-
-        selected = pills(
-            "Choose a question to get started or write your own below.",
-            [
-                "What is Snowflake?",
-                "What company did Snowflake announce they would acquire in October 2023?",
-                "What company did Snowflake acquire in March 2022?",
-                "When did Snowflake IPO?",
-            ],
-            clearable=True,
-            index=None,
-        )
-
-        if "chat_engine" not in st.session_state:  # Initialize the query engine
-            st.session_state["chat_engine"] = index.as_chat_engine(
-                chat_mode="context", verbose=True
+            with gr.Row():
+                chat_window = gr.Chatbot(
+                    label="Message History",
+                    scale=3,
+                )
+                console = gr.HTML(elem_id="box")
+            with gr.Row():
+                message = gr.Textbox(label="Write A Message", scale=4)
+                clear = gr.ClearButton()
+
+            message.submit(
+                self._handle_user_message,
+                [message, chat_window],
+                [message, chat_window],
+                queue=False,
+            ).then(
+                self._generate_response,
+                chat_window,
+                [chat_window, console],
             )
+            clear.click(self._reset_chat, None, [message, chat_window, console])
 
-        for message in st.session_state["messages"]:  # Display the prior chat messages
-            with st.chat_message(message["role"]):
-                st.write(message["content"])
-
-        # To avoid duplicated display of answered pill questions each rerun
-        if selected and selected not in st.session_state.get(
-            "displayed_pill_questions", set()
-        ):
-            st.session_state.setdefault("displayed_pill_questions", set()).add(selected)
-            with st.chat_message("user"):
-                st.write(selected)
-            with st.chat_message("assistant"):
-                response = st.session_state["chat_engine"].stream_chat(selected)
-                response_str = ""
-                response_container = st.empty()
-                for token in response.response_gen:
-                    response_str += token
-                    response_container.write(response_str)
-                add_to_message_history("user", selected)
-                add_to_message_history("assistant", response)
-
-        if prompt := st.chat_input(
-            "Your question"
-        ):  # Prompt for user input and save to chat history
-            add_to_message_history("user", prompt)
-
-            # Display the new question immediately after it is entered
-            with st.chat_message("user"):
-                st.write(prompt)
-
-            # If last message is not from assistant, generate a new response
-            # if st.session_state["messages"][-1]["role"] != "assistant":
-            with st.chat_message("assistant"):
-                response = st.session_state["chat_engine"].stream_chat(prompt)
-                response_str = ""
-                response_container = st.empty()
-                for token in response.response_gen:
-                    response_str += token
-                    response_container.write(response_str)
-                # st.write(response.response)
-                add_to_message_history("assistant", response.response)
-
-            # Save the state of the generator
-            st.session_state["response_gen"] = response.response_gen
+        demo.launch(server_name="0.0.0.0", server_port=8080)
 
 
 if __name__ == "__main__":
-    StreamlitChatPack(run_from_main=True).run()
+    GradioReActAgentPack(run_from_main=True).run()