Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import os
 from openai import OpenAI
-import os.path
-from datetime import datetime
+from datetime import datetime, timezone, timedelta
+import hashlib
+import hmac
 
 ################# Start PERSONA-SPECIFIC VALUES ######################
 coach_code = os.getenv("COACH_CODE")
@@ -23,24 +24,49 @@ tx = os.getenv("TX")
 prefix = os.getenv("PREFIX")  #  "/data/" if in HF or "data/" if local
 file_name = os.getenv("FILE_NAME")
 
+############### VERIFY USER ###################
+def generate_access_code(time):
+    secret = os.getenv("SHARED_SECRET_KEY")
+    
+    time_block = time.replace(minute=(time.minute // 10) * 10, second=0, microsecond=0)
+    time_string = time_block.strftime('%Y%m%d%H%M')
+    
+    hmac_obj = hmac.new(secret.encode(), time_string.encode(), hashlib.sha256)
+    hmac_digest = hmac_obj.hexdigest()
+    
+    xor_result = bytes(int(hmac_digest[i], 16) ^ int(hmac_digest[-4+i], 16) for i in range(4))
+    
+    return xor_result.hex()[:4]
+
+def verify_code(code, access_granted):
+    now = datetime.now(timezone.utc)
+    
+    codes = [generate_access_code(now + timedelta(minutes=offset)) 
+             for offset in [-20, -10, 0, 10, 20]]
+    
+    if code in codes:
+        return True, gr.update(interactive=True), gr.update(interactive=True), "Access granted. Please proceed to the Chat tab."
+    else:
+        return False, gr.update(interactive=False), gr.update(interactive=False), "Incorrect code. Please try again."
+
 ############### CHAT ###################
-def predict(user_input, history):
-    max_length = 3000
+def predict(user_input, history, access_granted):
+    if not access_granted:
+        return history, "Access not granted. Please enter the correct code in the Access tab."
+    
+    max_length = 1000
+    if len(user_input) > max_length:
+        user_input = ""
+
     transcript_file_path = f"{prefix}{coach_code}-{file_name}"
-    transcript = ""  # Initialize the transcript variable
 
     if user_input == tx + coach_code:
         try:
-            # Prepare the transcript for the Textbox output
             if os.path.exists(transcript_file_path):
                 with open(transcript_file_path, "r", encoding="UTF-8") as file:
-                    transcript = file.read()
-                return transcript
+                    return history, file.read()
         except FileNotFoundError:
-            return "File '" + file_name + "' not found."
-    elif len(user_input) > max_length:
-        user_input = ""
-        # raise gr.Error(f"Input is TOO LONG. Max length is {max_length} characters. Try again.")
+            return history, "File '" + file_name + "' not found."
 
     history_openai_format = [
         {"role": "system", "content": "IDENTITY: " + sys_prompt_new}
@@ -59,26 +85,81 @@ def predict(user_input, history):
         stream=True
     )
 
-    output_stream = ""
-    try:
-        for chunk in completion:
-            if chunk.choices[0].delta.content is not None:
-                output_stream = output_stream + (chunk.choices[0].delta.content)
-        message_content = output_stream
-    except StopAsyncIteration:
-        pass
+    message_content = ""
+    for chunk in completion:
+        if chunk.choices[0].delta.content is not None:
+            message_content += chunk.choices[0].delta.content
 
     # Append latest user and assistant messages to the transcript
-    transcript += "Date/Time: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n"
+    transcript = f"Date/Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
     transcript += f"YOU: {user_input}\n\n"
     transcript += f"{coach_name_upper}: {message_content}\n\n\n"
+    
     # Write the updated transcript to the file
     with open(transcript_file_path, "a", encoding="UTF-8") as file:
         file.write(transcript)
 
-    return message_content
+    history.append((user_input, message_content))
+    return history, ""
+
+with gr.Blocks(theme, css="""
+    #chatbot { flex-grow: 1; height: 460px; overflow-y: auto; }
+    .gradio-container { height: 700px; max-width: 100% !important; padding: 0 !important; }
+    #component-0 { height: 100%; }
+    #component-3 { height: calc(100% - 250px); }
+    footer { display: none !important; }
+    #submit-btn { margin-top: 10px; }
+    #code_submit {
+        height: 50px !important;
+        font-size: 1.2em !important;
+    }
+    #code_message { 
+        font-size: 1.4em !important; 
+        font-weight: bold !important; 
+        padding: 15px !important; 
+        margin-top: 15px !important;
+        border-radius: 5px !important;
+        text-align: center !important;
+    }
+    #code_message.error { color: #d32f2f !important; background-color: #ffcdd2 !important; }
+    #code_message.success { color: #388e3c !important; background-color: #c8e6c9 !important; }
+    @media (max-width: 600px) {
+        #code_submit {
+            height: 60px !important;
+            font-size: 1.3em !important;
+        }
+        #code_message {
+            font-size: 1.2em !important;
+            padding: 10px !important;
+        }
+    }
+""") as demo:
+    access_granted = gr.State(False)
+
+    with gr.Tab("Access"):
+        gr.Markdown("Enter the Access Code displayed in the upper-left corner.")
+        code_input = gr.Textbox(label="Access Code", type="text", placeholder="Enter CODE here...")
+        code_submit = gr.Button("Submit Code", elem_id="code_submit")
+        code_message = gr.Markdown(elem_id="code_message")
+
+    with gr.Tab("Chat"):
+        chatbot = gr.Chatbot(label="Conversation")
+        msg = gr.Textbox(
+            label=f"Chat with {coach_name_short}", 
+            placeholder="Type your message here... (MAX: 1000 characters)",
+            interactive=False
+        )
+        submit = gr.Button("Submit Message", interactive=False)
+
+    def submit_code(code, access_granted):
+        success, _, _, message = verify_code(code, access_granted)
+        css_class = "success" if success else "error"
+        return success, gr.update(interactive=success), gr.update(interactive=success), gr.update(value=message, elem_classes=[css_class])
+
+    code_input.submit(submit_code, inputs=[code_input, access_granted], outputs=[access_granted, msg, submit, code_message])
+    code_submit.click(submit_code, inputs=[code_input, access_granted], outputs=[access_granted, msg, submit, code_message])
+
+    msg.submit(predict, [msg, chatbot, access_granted], [chatbot, msg])
+    submit.click(predict, [msg, chatbot, access_granted], [chatbot, msg])
 
-#GUI
-with gr.Blocks(theme, css="footer {visibility: hidden}") as demo:
-    gr.ChatInterface(predict, submit_btn="Chat with "+ coach_name_short, retry_btn=None, undo_btn=None, clear_btn=None, autofocus=True)
 demo.launch(show_api=False)