Upload folder using huggingface_hub

.gitignore

@@ -4,6 +4,7 @@
 # Data
 chroma_db/
 data/
+sessions.pkl
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

app.py

@@ -7,6 +7,8 @@ from tools.squad_tools import SquadRetrieverTool, SquadQueryTool
 from tools.text_to_image import TextToImageTool
 from gradio.context import Context
 from gradio import Request
+import pickle
+import os
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -17,6 +19,11 @@ TASK_SOLVING_TOOLBOX = [
     TextToImageTool(),
 ]
 
+sessions_path = "sessions.pkl"
+sessions = pickle.load(open(sessions_path, "rb")) if os.path.exists(sessions_path) else {}
+
+app = None
+
 model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 # model_name = "http://localhost:1234/v1"
 
@@ -46,11 +53,10 @@ def add_message(message, messages):
     messages.append(ChatMessage(role="user", content=message))
     return messages
 
-sessions = {}
-
 def interact_with_agent(messages, request: Request):
+    username = request.username
     prompt = messages[-1]['content']
-    agent.logs = sessions.get(request.session_hash + "_logs", [])
+    agent.logs = sessions.get(username + "_logs", [])
     for msg in stream_from_transformers_agent(agent, prompt):
         messages.append(msg)
         yield messages
@@ -59,15 +65,18 @@ def interact_with_agent(messages, request: Request):
 def persist(component):
 
     def resume_session(value, request: Request):
-        print(f"Resuming session for {request.session_hash}")
-        state = sessions.get(request.session_hash, value)
-        agent.logs = sessions.get(request.session_hash + "_logs", [])
+        username = request.username
+        print(f"Resuming session for {username}")
+        state = sessions.get(username, value)
+        agent.logs = sessions.get(username + "_logs", [])
         return state
 
     def update_session(value, request: Request):
-        print(f"Updating persisted session state for {request.session_hash}")
-        sessions[request.session_hash] = value
-        sessions[request.session_hash + "_logs"] = agent.logs
+        username = request.username
+        print(f"Updating persisted session state for {username}")
+        sessions[username] = value
+        sessions[username + "_logs"] = agent.logs
+        pickle.dump(sessions, open(sessions_path, "wb"))
         return
 
     Context.root_block.load(resume_session, inputs=[component], outputs=component)
@@ -75,7 +84,16 @@ def persist(component):
 
     return component
 
+def welcome_message(request: Request):
+    return f"<h2>Welcome, {request.username}</h2>"
+
 with gr.Blocks(fill_height=True) as demo:
+    # put the welcome message and logout button in a row
+    with gr.Row() as row:
+        welcome_msg = gr.Markdown(f"Welcome")
+        logout_button = gr.Button("Logout", link="/logout")
+        demo.load(welcome_message, None, welcome_msg)
+
     chatbot = persist(gr.Chatbot(
         value=[],
         label="SQuAD Agent",
@@ -85,11 +103,12 @@ with gr.Blocks(fill_height=True) as demo:
             "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png",
         ),
         scale=1,
-        # bubble_full_width=False,
         autoscroll=True,
         show_copy_all_button=True,
         show_copy_button=True,
-        placeholder="Enter a message",
+        placeholder="""<h1>SQuAD Agent</h1>
+            <h2>I am your friendly guide to the Stanford Question and Answer Dataset (SQuAD).</h2>
+        """,
         examples=[
             {
                 "text": "What is on top of the Notre Dame building?",
@@ -110,5 +129,16 @@ with gr.Blocks(fill_height=True) as demo:
         interact_with_agent, [chatbot], [chatbot]
     )
 
+def honor_system_auth(username, password):
+    return password == "happy"
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        auth=honor_system_auth,
+        auth_message="""<h3>Honor System Authentication:</h3>
+            <p>Log in with <strong>any username</strong> and the password <strong>"happy"</strong>.</p>
+            <br/><br/>
+            <i>Note:Your chat history will be saved to your username, but take care because others 
+            may log in with the same username and see your chat history.</i>
+        """,
+    )