Persist sessions through restarts with sqlite

Author: deckar01

.gitignore

@@ -1,6 +1,7 @@
 .venv/
 .offload/
 .history/
+oracle.db
 
 __pycache__
 .DS_Store

oracle/__init__.py

@@ -1,3 +1,5 @@
+DB_PATH = 'oracle.db'
+
 def get_device():
     import torch
 

oracle/gradio/__main__.py

@@ -1,17 +1,37 @@
 import os
 import hashlib
-from oracle.gradio.gui import demo
+import sqlite3
 
+from oracle import DB_PATH
+from oracle.gradio.gui import demo
 
-credentials = {}
 
 def open_registration(username, password):
+    connection = sqlite3.connect(DB_PATH)
     password_hash = hashlib.sha256(password.encode()).hexdigest()
-    saved_hash = credentials.get(username, password_hash)
-    credentials[username] = saved_hash
+    query = connection.execute(
+        "SELECT password FROM credentials WHERE user=?",
+        (username,),
+    )
+    saved_hash = query.fetchone()
+    if not saved_hash:
+        connection.execute(
+            "INSERT OR REPLACE INTO credentials VALUES (?, ?)",
+            (username, password_hash),
+        )
+        connection.commit()
+        saved_hash = password_hash
+    else:
+        saved_hash = saved_hash[0]
+    connection.close()
     return password_hash == saved_hash
 
 if __name__ == '__main__':
+    connection = sqlite3.connect(DB_PATH)
+    try: connection.execute("CREATE TABLE credentials (user TEXT PRIMARY KEY, password TEXT)")
+    except: pass
+    connection.close()
+
     os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
     demo.launch(
         server_port=8080,

oracle/gradio/gui.py

@@ -9,20 +9,20 @@
     # Model
 
     session_state = gr.State(ChatSession)
-    raw_chat_log = persist(gr.State(list))
+    raw_chat_log = persist("chat", gr.State(list))
 
     # View
 
     with gr.Accordion('☰', open=False, elem_id='settings'):
         with gr.Tab('Context'):
-            context_input = persist(gr.Dropdown(label='Source'))
-            motive_input = persist(gr.Textbox(label='Motivation'))
-            keyword_checkbox = persist(gr.Checkbox(True, label='Ask the model for keywords?'))
-            debug_checkbox = persist(gr.Checkbox(label='Show debug info?'))
+            context_input = persist("context", gr.Dropdown(label='Source'))
+            motive_input = persist("motive", gr.Textbox(label='Motivation'))
+            keyword_checkbox = persist("keyword", gr.Checkbox(True, label='Ask the model for keywords?'))
+            debug_checkbox = persist("debug", gr.Checkbox(label='Show debug info?'))
 
         with gr.Tab('Model'):
-            model_input = persist(gr.Dropdown(label='Chat Model'))
-            style_input = persist(gr.Dropdown(
+            model_input = persist("model", gr.Dropdown(label='Chat Model'))
+            style_input = persist("style", gr.Dropdown(
                 label='Response Style',
                 choices=STYLES,
                 allow_custom_value=True,

oracle/gradio/utils.py

@@ -1,10 +1,13 @@
 import inspect
-from uuid import uuid4
+import sqlite3
+import pickle
 
 import gradio as gr
 from gradio.context import Context
 from gradio.events import Changeable
 
+from oracle import DB_PATH
+
 
 def locked(**kwargs):
     return gr.update(**kwargs, interactive=False)
@@ -26,23 +29,31 @@ def wrapper(fn):
     else:
         return wrapper
 
-user_sessions = {}
-
-def get_session_id(request):
-    session_id = user_sessions.get(request.username, None)
-    if not session_id:
-        session_id = uuid4()
-    if request.username:
-        user_sessions[request.username] = session_id
-    return session_id
+def persist(name, component):
+    connection = sqlite3.connect(DB_PATH)
+    try: connection.execute(f"CREATE TABLE {name} (user TEXT PRIMARY KEY, value TEXT)")
+    except: pass
+    connection.close()
 
-def persist(component):
-    sessions = {}
     @on(Context.root_block.load)
     def resume_session(value: component, request: gr.Request) -> component:
-        return sessions.get(get_session_id(request), value)
+        connection = sqlite3.connect(DB_PATH)
+        query = connection.execute(
+            f"SELECT value FROM {name} WHERE user=?",
+            (request.username,),
+        )
+        saved_value = query.fetchone()
+        connection.close()
+        return pickle.loads(saved_value[0]) if saved_value else value
+
     def update_session(value: component, request: gr.Request):
-        sessions[get_session_id(request)] = value
+        connection = sqlite3.connect(DB_PATH)
+        query = connection.execute(
+            f"INSERT OR REPLACE INTO {name} VALUES (?, ?)",
+            (request.username, pickle.dumps(value, 5)),
+        )
+        connection.commit()
+        connection.close()
 
     if hasattr(component, 'change'):
         on(component.change, update_session)