handle multiple logins

app.py

@@ -5,9 +5,11 @@ from tools import MedicalTools
 from graph_setup import GraphSetup
 from chat_handler import ChatHandler
 from audio_handler import AudioHandler
-
+from user_data import initialize_db, add_user, create_session
+import uuid
 
 load_dotenv(override=True)
+initialize_db()
 
 rag = RAG_Setup()
 medical_tools = MedicalTools(rag)
@@ -18,20 +20,75 @@ chat_handler = ChatHandler(graph, rag)
 audio_handler = AudioHandler()
 
 
-def transcribe_audio_wrapper(audio, current_text, file_input, message_history):
+def transcribe_audio_wrapper(audio, current_text, file_input, message_history, user_state, session_state):
+    if not user_state or not session_state:
+        warning = {
+            "role": "assistant",
+            "content": "Please log in and start a session before using voice input."
+        }
+        return message_history + [warning], current_text, None, file_input, user_state, session_state
+
     return audio_handler.transcribe_audio(
         audio, 
         current_text, 
         file_input, 
         message_history, 
+        user_state,
+        session_state,
         chat_handler.chat
     )
 
+def handle_login(user_identifier):
+    if not user_identifier or not user_identifier.strip():
+        return (
+            "Please enter a user name or email.",
+            None,
+            None
+        )
+
+    user_id = user_identifier.strip().lower()
+    session_id = str(uuid.uuid4())
+
+    add_user(user_id, user_identifier)
+    create_session(user_id, session_id)
+
+    session_md = f"**Active user:** {user_id}<br>**Session:** {session_id}"
+
+    return (
+        session_md,
+        {"user_id": user_id, "name": user_identifier},
+        {"session_id": session_id}
+    )
+
+
+def handle_logout():
+    return "No active session.", None, None
 
 with gr.Blocks(title="Medical Assistant") as demo:
+    user_state = gr.State(value=None)
+    session_state = gr.State(value=None)
+
     gr.Markdown("# 🏥 Medical Assistant")
     gr.Markdown("Ask questions using text, voice, or upload medical documents")
     
+    with gr.Accordion("User Login", open=True):
+        user_input = gr.Textbox(label="Enter email or username", placeholder="name@example.com")
+        with gr.Row():
+            login_button = gr.Button("Start Session", variant="primary")
+            logout_button = gr.Button("End Session", variant="stop")
+        session_display = gr.Markdown("No active session.")
+    
+    login_button.click(
+        handle_login,
+        inputs=[user_input],
+        outputs=[session_display, user_state, session_state],
+    )
+
+    logout_button.click(
+        handle_logout,
+        outputs=[session_display, user_state, session_state],
+    )
+
     chatbot = gr.Chatbot(label="Conversation", height=400)
     
     with gr.Row():
@@ -62,22 +119,21 @@ with gr.Blocks(title="Medical Assistant") as demo:
     
     submit_btn.click(
         chat_handler.chat,
-        inputs=[text_input, file_input, chatbot],
-        outputs=[chatbot, text_input, file_input]
+        inputs=[text_input, file_input, chatbot, user_state, session_state],
+        outputs=[chatbot, text_input, file_input],
     )
     
     text_input.submit(
         chat_handler.chat,
-        inputs=[text_input, file_input, chatbot],
+        inputs=[text_input, file_input, chatbot, user_state, session_state],
         outputs=[chatbot, text_input, file_input]
     )
 
     audio_input.change(
         transcribe_audio_wrapper,
-        inputs=[audio_input, text_input, file_input, chatbot],
+        inputs=[audio_input, text_input, file_input, chatbot, user_state, session_state],
         outputs=[chatbot, text_input, audio_input, file_input] 
     )
 
-
 if __name__ == "__main__":
     demo.launch(share=True)

audio_handler.py

@@ -5,16 +5,18 @@ class AudioHandler:
     def __init__(self):
         self.transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-small")
     
-    def transcribe_audio(self, audio, current_text, file_input, message_history, chat_func):
+    def transcribe_audio(self, audio, current_text, file_input, message_history, user_state, session_state, chat_func):
         if audio is None:
-            return message_history, current_text, None, file_input
+            return message_history, current_text, None, file_input, user_state, session_state
         
         transcript = self.transcriber(audio)["text"].strip()
         
-        updated_history, cleared_text, cleared_file = chat_func(
+        updated_history, cleared_text, cleared_file, user_state, session_state = chat_func(
             transcript, 
             file_input, 
-            message_history
+            message_history,
+            user_state,
+            session_state
         )
         
-        return updated_history, current_text, None, cleared_file
+        return updated_history, cleared_text, None, cleared_file, user_state, session_state

chat_handler.py

@@ -11,7 +11,14 @@ class ChatHandler:
         self.session_id = str(uuid.uuid4())
         print(self.session_id)
     
-    def chat(self, user_message, uploaded_file, message_history):
+    def chat(self, user_message, uploaded_file, message_history, user_state, session_state):
+        if not user_state or not session_state:
+            warning = {
+                "role": "assistant",
+                "content": "Please log in and start a session before chatting."
+            }
+            return message_history + [warning], user_message, uploaded_file, user_state, session_state
+        
         user_query_parts = []
         try: 
             if user_message and user_message.strip():

requirements.txt

@@ -6,4 +6,5 @@ langgraph-checkpoint-sqlite
 langchain
 transformers 
 sentence-transformers
-torch
+torch
+pypdf

user_data.py

@@ -0,0 +1,88 @@
+import sqlite3
+from pathlib import Path
+
+DB_PATH = Path("data/user_data.db")
+
+def get_connection():
+    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(DB_PATH)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+def initialize_db():
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.executescript("""
+        CREATE TABLE IF NOT EXISTS users (
+            id TEXT PRIMARY KEY,
+            name TEXT NOT NULL,
+            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP                 
+        );
+                         
+        CREATE TABLE IF NOT EXISTS sessions (
+            id TEXT PRIMARY KEY,
+            user_id TEXT NOT NULL,
+            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+            FOREIGN KEY (user_id) REFERENCES users(id)
+        );
+
+        CREATE TABLE IF NOT EXISTS document_classifications (
+            file_hash TEXT PRIMARY KEY,
+            doc_type TEXT NOT NULL,
+            classified_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+        );
+    """)
+    conn.commit()
+    conn.close()
+
+def add_user(user_id, name):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("""INSERT INTO users (id, name) VALUES (?, ?)
+                   ON CONFLICT(id) DO UPDATE SET name = excluded.name""", (user_id, name))
+    conn.commit()
+    conn.close()
+
+def create_session(user_id, session_id):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("INSERT INTO sessions (id, user_id) VALUES (?, ?)", (session_id, user_id))
+    conn.commit()
+    conn.close()
+
+def user_exists(user_id):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("SELECT 1 FROM users WHERE id = ?", (user_id,))
+    exists = cursor.fetchone() is not None
+    conn.close()
+    return exists
+
+
+def get_document_label(file_hash: str):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute(
+        "SELECT doc_type FROM document_classifications WHERE file_hash = ?",
+        (file_hash,)
+    )
+    row = cursor.fetchone()
+    conn.close()
+    return row["doc_type"] if row else None
+
+
+def save_document_label(file_hash: str, doc_type: str):
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute(
+        """
+        INSERT INTO document_classifications (file_hash, doc_type)
+        VALUES (?, ?)
+        ON CONFLICT(file_hash) DO UPDATE SET
+            doc_type = excluded.doc_type,
+            classified_at = CURRENT_TIMESTAMP
+        """,
+        (file_hash, doc_type)
+    )
+    conn.commit()
+    conn.close()