easier session id logic

app.py

@@ -5,82 +5,83 @@ import uuid
 import os
 import json
 
-# ---- Constants ----
+# ---- Configuration ----
 MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
 DATASET_REPO = "frimelle/companion-chat-logs"
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# ---- System Prompt ----
+# ---- Load system prompt ----
 with open("system_prompt.txt", "r") as f:
     SYSTEM_PROMPT = f.read()
 
 client = InferenceClient(MODEL_NAME)
 api = HfApi()
 
-# ---- Log function (writes per session) ----
-def append_to_session_log(user_message, assistant_message, session_id):
-    log_entry = {
+# ---- Session ID and local file setup ----
+SESSION_ID = 0
+LOCAL_LOG_PATH = f"chatlog_{SESSION_ID}.jsonl"
+REMOTE_LOG_PATH = f"sessions/{SESSION_ID}.jsonl"  # stored in a folder in the dataset
+
+def increment_session_id():
+    global SESSION_ID
+    SESSION_ID += 1
+    return SESSION_ID
+
+# ---- Logging per session ----
+def append_to_session_log(user_message, assistant_message):
+    row = {
         "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")
 
-    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
+    # Push to HF dataset
     api.upload_file(
-        path_or_fileobj=local_path,
-        path_in_repo=remote_path,
+        path_or_fileobj=LOCAL_LOG_PATH,
+        path_in_repo=REMOTE_LOG_PATH,
         repo_id=DATASET_REPO,
         repo_type="dataset",
-        token=HF_TOKEN,
+        token=HF_TOKEN
     )
 
-# ---- Chatbot logic ----
-def respond(message, history, session_id):
+# ---- Chatbot function ----
+def respond(message, history):
     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 = ""
+
     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
 
-    # 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
+    # Save after each message pair
+    append_to_session_log(message, response)
+    increment_session_id()
 
-    msg.submit(user_submit, inputs=[msg, chatbot, session_id], outputs=[chatbot, session_id])
-    clear.click(lambda: None, None, chatbot)
+# ---- Gradio Interface ----
+demo = gr.ChatInterface(
+    respond,
+    title="BoundrAI",
+)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()