deploy at 2024-08-24 15:14:26.631580

main.py

@@ -57,8 +57,9 @@ import tempfile
 from enum import Enum
 from typing import Tuple as T
 from urllib.parse import quote
+import uuid
 
-DEV_MODE = False
+DEV_MODE = True
 
 if DEV_MODE:
     print("Running in DEV_MODE - Hot reload enabled")
@@ -170,7 +171,7 @@ middlewares = [
         SessionMiddleware,
         secret_key=get_key(fname=sess_key_path),
         max_age=3600,
-        same_site="None",
+        same_site='lax',  # This allows cookies to be sent in top-level navigations
     ),
     Middleware(XFrameOptionsMiddleware),
 ]
@@ -202,8 +203,7 @@ app, rt = fast_app(
 
 
 sesskey = get_key(fname=sess_key_path)
-if DEV_MODE:
-    print(f"Session key: {sesskey}")
+print(f"Session key: {sesskey}")
 
 
 # enum class for rank profiles
@@ -299,6 +299,7 @@ def get(sess):
         "Treating Asthma With Plants vs. Pills",
         "Alkylphenol Endocrine Disruptors",
         "Testing Turmeric on Smokers",
+        "The Role of Pesticides in Parkinson's Disease",
     ]
     return (
         Title("Vespa demo"),
@@ -453,9 +454,23 @@ def replace_hi_with_strong(text):
 
 
 def log_query_to_db(query, ranking, sess):
-    return queries.insert(
-        Query(query=query, ranking=ranking, sess_id=sesskey, timestamp=int(time.time()))
-    )
+    if 'user_id' not in sess:
+        sess['user_id'] = str(uuid.uuid4())
+    
+    if 'queries' not in sess:
+        sess['queries'] = []
+    
+    query_data = {
+        'query': query,
+        'ranking': ranking,
+        'timestamp': int(time.time())
+    }
+    sess['queries'].append(query_data)
+    
+    # Limit the number of queries stored in the session to prevent it from growing too large
+    sess['queries'] = sess['queries'][-100:]  # Keep only the last 100 queries
+
+    return query_data
 
 
 def parse_results(records):
@@ -544,15 +559,8 @@ def get_yql(ranking: RankProfile, userquery: str) -> T[str, dict]:
 
 @app.get("/search")
 async def search(userquery: str, ranking: str, sess):
-    print("In search")
-    print(f"Session: {sess}")
-    if "queries" not in sess:
-        print("Creating queries list")
-        sess["queries"] = []
+    print(sess)
     quoted = quote(userquery) + "&ranking=" + ranking
-    sess["queries"].append(quoted)
-    print(f"Searching for: {userquery}")
-    print(f"Ranking: {ranking}")
     log_query_to_db(userquery, ranking, sess)
     yql, body = get_yql(ranking, userquery)
     async with vespa_app.asyncio() as session:
@@ -593,16 +601,23 @@ async def search(userquery: str, ranking: str, sess):
 
 
 @app.get("/download_csv")
-def download_csv(auth):
-    queries_dict = list(db.query("SELECT * FROM queries"))
-    queries = [Query(**query) for query in queries_dict]
+def download_csv(auth, sess):
+    if not auth:
+        return RedirectResponse("/login", status_code=303)
+
+    all_queries = []
+    for user_id, user_sess in app.state.sessions.items():
+        user_queries = user_sess.get('queries', [])
+        for query in user_queries:
+            query['user_id'] = user_id
+        all_queries.extend(user_queries)
 
     # Create CSV in memory
     csv_file = StringIO()
     csv_writer = csv.writer(csv_file)
-    csv_writer.writerow(["Query", "Session ID", "Timestamp"])
-    for query in queries:
-        csv_writer.writerow([query.query, query.sess_id, query.timestamp])
+    csv_writer.writerow(["Query", "User ID", "Timestamp", "Ranking"])
+    for query in all_queries:
+        csv_writer.writerow([query['query'], query['user_id'], query['timestamp'], query['ranking']])
 
     # Move to the beginning of the StringIO object
     csv_file.seek(0)
@@ -621,21 +636,27 @@ def download_csv(auth):
 
 
 @app.get("/admin")
-def get_admin(auth, page: int = 1):
+def get_admin(auth, sess, page: int = 1):
+    if not auth:
+        return RedirectResponse("/login", status_code=303)
+
     limit = 15
     offset = (page - 1) * limit
-    total_queries_result = list(
-        db.query("SELECT COUNT(*) AS count FROM queries ORDER BY timestamp DESC")
-    )
-    total_queries = total_queries_result[0]["count"]
-    queries_dict = list(
-        db.query(f"SELECT * FROM queries LIMIT {limit} OFFSET {offset}")
-    )
-    queries = [Query(**query) for query in queries_dict]
 
-    total_pages = (
-        total_queries + limit - 1
-    ) // limit  # Calculate total number of pages
+    all_queries = []
+    for user_id, user_sess in app.state.sessions.items():
+        user_queries = user_sess.get('queries', [])
+        for query in user_queries:
+            query['user_id'] = user_id
+        all_queries.extend(user_queries)
+
+    # Sort queries by timestamp in descending order
+    all_queries.sort(key=lambda x: x['timestamp'], reverse=True)
+
+    total_queries = len(all_queries)
+    queries = all_queries[offset:offset+limit]
+
+    total_pages = (total_queries + limit - 1) // limit
 
     # Define the range of pages to display
     page_window = 5  # Number of pages to display at once
@@ -715,16 +736,16 @@ def get_admin(auth, page: int = 1):
                 Thead(
                     Tr(
                         Th("Query"),
-                        Th("Session ID"),
+                        Th("User ID"),
                         Th("Datetime"),
                     )
                 ),
                 Tbody(
                     *[
                         Tr(
-                            Td(query.query),
-                            Td(query.sess_id),
-                            Td(query.get_datetime()),
+                            Td(query['query']),
+                            Td(query['user_id']),
+                            Td(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(query['timestamp']))),
                         )
                         for query in queries
                     ],
@@ -806,8 +827,6 @@ def get_about(auth, sess):
 
 @app.get("/document/{docid}")
 def get_document(docid: str, sess):
-    print(f"Getting document {docid}")
-    print(f"Session: {sess}")
     resp = vespa_app.get_data(data_id=docid, schema="doc", namespace="tutorial")
     doc = resp.json
     # Link with Back to search results at top of page
@@ -831,4 +850,17 @@ if not DEV_MODE:
         setup_hf_backup(app)
     except Exception as e:
         print(f"Error setting up hf backup: {e}")
+
+# Add this at the end of your file, before serve()
+app.state.sessions = {}
+
+@app.middleware("http")
+async def add_session_to_app_state(request, call_next):
+    response = await call_next(request)
+    if 'session' in request.scope:
+        session = request.scope['session']
+        if 'user_id' in session:
+            app.state.sessions[session['user_id']] = session
+    return response
+
 serve()