[#122] Replace cookie with local storage for token (#123)

app.py

@@ -80,7 +80,8 @@ css = """
 """
 
 with gr.Blocks(title="Yanolja Arena", css=css) as app:
-  set_token(app)
+  token = gr.Textbox(visible=False)
+  set_token(app, token)
 
   with gr.Row():
     gr.HTML("""
@@ -188,7 +189,7 @@ with gr.Blocks(title="Yanolja Arena", css=css) as app:
       ]).then(fn=get_responses,
               inputs=[
                   prompt_textarea, category_radio, source_language,
-                  target_language
+                  target_language, token
               ],
               outputs=response_boxes + model_names + [instruction_state])
   submit_event.success(fn=lambda: gr.Row(visible=True), outputs=vote_row)

rate_limit.py

@@ -55,7 +55,7 @@ class RateLimiter:
     self.scheduler.start()
 
   def check_rate_limit(self, token: str):
-    if not token or token not in self.last_request_times:
+    if not token or not self.token_exists(token):
       raise InvalidTokenException()
 
     if (datetime.now() - self.last_request_times[token]).seconds < 5:
@@ -70,6 +70,9 @@ class RateLimiter:
   def initialize_request(self, token: str):
     self.last_request_times[token] = datetime.min
 
+  def token_exists(self, token: str):
+    return token in self.last_request_times
+
   def _remove_old_tokens(self):
     for token, last_request_time in dict(self.last_request_times).items():
       if (datetime.now() - last_request_time).days >= 1:
@@ -82,23 +85,33 @@ class RateLimiter:
 rate_limiter = RateLimiter()
 
 
-def set_token(app: gr.Blocks):
+def set_token(app: gr.Blocks, token: gr.Textbox):
+
+  get_client_token = """
+  function() {
+    return localStorage.getItem("arena_token");
+  }
+  """
+
+  def set_server_token(existing_token):
+    if existing_token and rate_limiter.token_exists(existing_token):
+      return existing_token
 
-  def set_token_server():
     new_token = uuid4().hex
     rate_limiter.initialize_request(new_token)
     return new_token
 
-  set_token_client = """
+  set_client_token = """
   function(newToken) {
-    const expiresDateString = new Date(Date.now() + 24 * 60 * 60 * 1000).toUTCString();
-    document.cookie = `arena_token=${newToken}; expires=${expiresDateString};`;
+    localStorage.setItem("arena_token", newToken);
   }
   """
 
-  token = gr.Textbox(visible=False)
-  app.load(fn=set_token_server, outputs=[token])
-  token.change(fn=lambda _: None, js=set_token_client, inputs=[token])
+  app.load(fn=set_server_token,
+           js=get_client_token,
+           inputs=[token],
+           outputs=[token])
+  token.change(fn=lambda _: None, js=set_client_token, inputs=[token])
 
 
 def signal_handler(sig, frame):

response.py

@@ -3,7 +3,6 @@ This module contains functions for generating responses using LLMs.
 """
 
 import enum
-from http import cookies
 import logging
 from random import sample
 from typing import List
@@ -66,7 +65,7 @@ def get_instruction(category: str, model: Model, source_lang: str,
 
 
 def get_responses(prompt: str, category: str, source_lang: str,
-                  target_lang: str, request: gr.Request):
+                  target_lang: str, token: str):
   if not category:
     raise gr.Error("Please select a category.")
 
@@ -74,10 +73,6 @@ def get_responses(prompt: str, category: str, source_lang: str,
                                                not target_lang):
     raise gr.Error("Please select source and target languages.")
 
-  cookie = cookies.SimpleCookie()
-  cookie.load(request.headers["cookie"])
-  token = cookie["arena_token"].value if "arena_token" in cookie else None
-
   try:
     rate_limiter.check_rate_limit(token)
   except rate_limit.InvalidTokenException as e: