Update app.py

app.py

@@ -2,7 +2,7 @@ from langchain.agents import AgentExecutor, AgentType, initialize_agent
 from langchain.agents.structured_chat.prompt import SUFFIX
 from langchain.chat_models import ChatOpenAI
 from langchain.memory import ConversationBufferMemory
-from tools import generate_image_tool
+from tools import rag
 
 import chainlit as cl
 from chainlit.action import Action
@@ -19,7 +19,7 @@ def rename(orig_author):
     mapping = {
         "AgentExecutor": "The LLM Brain",
         "LLMChain": "The Assistant",
-        "GenerateImage": "DALL-E 3",
+        "RAG": "Jonah",
         "ChatOpenAI": "GPT-4 Turbo",
         "Chatbot": "Coolest App",
     }
@@ -81,13 +81,12 @@ async def setup_agent(settings):
     # This suffix is used to provide the chat history to the prompt.
     _SUFFIX = "Chat history:\n{chat_history}\n\n" + SUFFIX
 
-    # We initialize our agent here, which is simply being used to decide between responding with text
-    # or an image
+    # We initialize our agent here, which is simply being used to decide between responding with llm or tool
     agent = initialize_agent(
         llm=llm,  # our LLM (default is GPT-4 Turbo)
         tools=[
-            generate_image_tool
-        ],  # our custom tool used to generate images with DALL-E 3
+            rag
+        ],  # our custom tool used to retrieve context
         agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,  # the agent type we're using today
         memory=memory,  # our memory!
         agent_kwargs={
@@ -114,7 +113,7 @@ async def main(message: cl.Message):
     it back to the user.
     """
     agent = cl.user_session.get("agent")
-    cl.user_session.set("generated_image", None)
+    cl.user_session.set("rag", None)
 
     res = await cl.make_async(agent.run)(
         input=message.content, callbacks=[cl.LangchainCallbackHandler()]
@@ -123,15 +122,10 @@ async def main(message: cl.Message):
     elements = []
     actions = []
 
-    generated_image_name = cl.user_session.get("generated_image")
-    generated_image = cl.user_session.get(generated_image_name)
-    if generated_image:
+    tool_res = cl.user_session.get("rag")
+    if tool_res:
         elements = [
-            cl.Image(
-                content=generated_image,
-                name=generated_image_name,
-                display="inline",
-            )
+            tool_res
         ]
 
-    await cl.Message(content=res, elements=elements, actions=actions).send()
+    await cl.Message(content=res, elements=elements, actions=actions).send()