improve session id

app.py

@@ -5,77 +5,82 @@ import uuid
 import os
 import json
 
-# ---- Configuration ----
+# ---- Constants ----
 MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
 DATASET_REPO = "frimelle/companion-chat-logs"
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# ---- Load system prompt ----
+# ---- System Prompt ----
 with open("system_prompt.txt", "r") as f:
     SYSTEM_PROMPT = f.read()
 
 client = InferenceClient(MODEL_NAME)
 api = HfApi()
 
-# ---- Session ID and local file setup ----
-SESSION_ID = str(uuid.uuid4())
-LOCAL_LOG_PATH = f"chatlog_{SESSION_ID}.jsonl"
-REMOTE_LOG_PATH = f"sessions/{SESSION_ID}.jsonl"  # stored in a folder in the dataset
-
-# ---- Logging per session ----
-def append_to_session_log(user_message, assistant_message):
-    row = {
+# ---- Log function (writes per session) ----
+def append_to_session_log(user_message, assistant_message, session_id):
+    log_entry = {
         "timestamp": datetime.now().isoformat(),
         "user": user_message,
         "assistant": assistant_message,
         "system_prompt": SYSTEM_PROMPT,
-        "session_id": SESSION_ID
+        "session_id": session_id,
     }
-    with open(LOCAL_LOG_PATH, "a", encoding="utf-8") as f:
-        f.write(json.dumps(row) + "\n")
 
-    # Push to HF dataset
+    local_path = f"chatlog_{session_id}.jsonl"
+    remote_path = f"sessions/{session_id}.jsonl"
+
+    # Append to local file
+    with open(local_path, "a", encoding="utf-8") as f:
+        f.write(json.dumps(log_entry) + "\n")
+
+    # Push updated file to the dataset
     api.upload_file(
-        path_or_fileobj=LOCAL_LOG_PATH,
-        path_in_repo=REMOTE_LOG_PATH,
+        path_or_fileobj=local_path,
+        path_in_repo=remote_path,
         repo_id=DATASET_REPO,
         repo_type="dataset",
-        token=HF_TOKEN
+        token=HF_TOKEN,
     )
 
-# ---- Chatbot function ----
-def respond(message, history):
+# ---- Chatbot logic ----
+def respond(message, history, session_id):
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    
     for user_msg, bot_msg in history:
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
         if bot_msg:
             messages.append({"role": "assistant", "content": bot_msg})
-    
     messages.append({"role": "user", "content": message})
-    response = ""
 
+    response = ""
     for chunk in client.chat_completion(
         messages,
         max_tokens=512,
-        stream=True,
         temperature=0.7,
         top_p=0.95,
+        stream=True,
     ):
         token = chunk.choices[0].delta.content
         if token:
             response += token
             yield response
 
-    # Save after each message pair
-    append_to_session_log(message, response)
+    # Append message pair to the user's session log
+    append_to_session_log(message, response, session_id)
+
+# ---- Gradio App ----
+with gr.Blocks() as demo:
+    session_id = gr.State(str(uuid.uuid4()))  # Unique ID per user
+
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(placeholder="Ask me anything...")
+    clear = gr.Button("Clear chat")
+
+    def user_submit(message, chat_history, session_id):
+        return respond(message, chat_history, session_id), session_id
 
-# ---- Gradio Interface ----
-demo = gr.ChatInterface(
-    respond,
-    title="Zephyr Chatbot",
-)
+    msg.submit(user_submit, inputs=[msg, chatbot, session_id], outputs=[chatbot, session_id])
+    clear.click(lambda: None, None, chatbot)
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()