Update interface.py

interface.py

@@ -1,8 +1,9 @@
 import re, uuid
 import base64
-
+import os
 import bcrypt
 import gradio as gr
+import boto3
 from gradio_pdf import PDF
 from pathlib import Path
 import time
@@ -14,6 +15,8 @@ from fpdf import FPDF
 REPORT_DIR = Path("reports")
 REPORT_DIR.mkdir(exist_ok=True)
 SALT = b'$2b$12$MC7djiqmIR7154Syul5Wme'
+s3_client = boto3.client("s3", region_name="ca-central-1", aws_access_key_id=os.getenv("AWS_ACCESS_KEY"), aws_secret_access_key=os.getenv("AWS_SECRET_KEY"))
+BUCKET_NAME = "molx-data-storage"
 
 USERS = {
     'test_user': b'$2b$12$MC7djiqmIR7154Syul5WmeQwebwsNOK5svMX08zMYhvpF9P9IVXe6',
@@ -100,7 +103,7 @@ class ChatInterface:
         return history, gr.Textbox(value=message, interactive=False)
 
     async def process_message(
-        self, message: str, display_image: Optional[str], chat_history: List[ChatMessage]
+        self, message: str, display_image: Optional[str], chat_history: List[ChatMessage], session_details: dict
     ) -> AsyncGenerator[Tuple[List[ChatMessage], Optional[str], str], None]:
         """
         Process a message and generate responses.
@@ -197,6 +200,55 @@ class ChatInterface:
             )
             yield chat_history, self.display_file_path
 
+        finally:
+            store_chat_history(session_details['username'], session_details['session_id'], chat_history)
+
+def store_chat_history(username, session_id, chat_history):
+    """
+    Store the chat history (Agent responses) in S3 as a text file with a unique name.
+
+    Args:
+        username (str): The username of the user.
+        session_id (str): A unique session identifier.
+        chat_history (list): A list of agent responses to be saved in the text file.
+
+    Returns:
+        str: The URL of the uploaded chat history in S3.
+    """
+    chat_history_text = "\n".join(chat_history)
+    timestamp = str(int(time.time()))  # Get current timestamp for unique name
+    chat_history_path = f"/tmp/{session_id}_chat_history_{timestamp}.txt"
+
+    with open(chat_history_path, "w") as f:
+        f.write(chat_history_text)
+
+    # Upload chat history to S3 with the timestamp
+    return upload_to_s3(chat_history_path, username, session_id)
+
+
+def upload_to_s3(file_path, username, session_id):
+    """
+    Upload a file to S3 under the user's folder and session ID.
+
+    Args:
+        file_path (str): The path to the file to upload.
+        username (str): The username of the user.
+        session_id (str): A unique session identifier.
+        file_type (str): The type of file being uploaded (image, report, etc.).
+
+    Returns:
+        str: The URL of the uploaded file in S3.
+    """
+    # Define the S3 object key
+    file_name = Path(file_path).name
+    s3_key = f"{username}/{session_id}/{file_name}"
+
+    # Upload the file to S3
+    s3_client.upload_file(file_path, BUCKET_NAME, s3_key)
+
+    # Generate the file URL
+    file_url = f"https://{BUCKET_NAME}.s3.amazonaws.com/{s3_key}"
+    return file_url
 
 def create_demo(agent, tools_dict):
     """
@@ -210,6 +262,7 @@ def create_demo(agent, tools_dict):
         gr.Blocks: Gradio Blocks interface
     """
     interface = ChatInterface(agent, tools_dict)
+    session_details = {}
 
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         auth_state = gr.State(False)
@@ -257,15 +310,6 @@ def create_demo(agent, tools_dict):
                             image_display = gr.Image(
                                 label="Image", type="filepath", height=685, container=True
                             )
-                            # with gr.Row():
-                            #     upload_button = gr.UploadButton(
-                            #         "📎 Upload X-Ray",
-                            #         file_types=["image"],
-                            #     )
-                            #     dicom_upload = gr.UploadButton(
-                            #         "📄 Upload DICOM",
-                            #         file_types=["file"],
-                            #     )
                             with gr.Row():
                                 analyze_btn = gr.Button("Analyze 1")
                                 analyze2_btn = gr.Button("Analyze 2")
@@ -276,39 +320,26 @@ def create_demo(agent, tools_dict):
                         
                         with gr.Tab(label="Report section"):                               
                             generate_report_btn = gr.Button("Generate Report")
-                            # diseases_df = gr.Dataframe(
-                            #                         headers=["Disease", "Info"],
-                            #                         datatype=["str", "str"],
-                            #                         interactive=False, visible=False, max_height=220)
                             conclusion_tb = gr.Textbox(label="Conclusion", interactive=False)
                             with gr.Row():
                                 approve_btn = gr.Button("Approve", visible=False)
-                                # reject_btn = gr.Button("Reject", visible=False)
                             download_pdf_btn  = gr.DownloadButton(label="📥 Download PDF", visible=False)
-                            # pdf_preview  = gr.HTML(visible=False)
-                            # pdf_preview  = gr.File(visible=False)
                             pdf_preview  = PDF(visible=False)
-                            # rejection_text   = gr.Textbox(
-                            #     show_label=False,
-                            #     visible=False,
-                            #     placeholder="Tell us what is wrong with the report",
-                            #     container=False,
-                            #     interactive=True
-                            # )
-                            # with gr.Row():
-                            #     submit_reject_btn = gr.Button("Submit", visible=False)
-                            #     cancel_reject_btn = gr.Button("Cancel", visible=False)
 
         # Event handlers
         def authenticate(username, password):
             hashed = USERS.get(username)
             if hashed and bcrypt.checkpw(password.encode(), hashed):
+                session_details["username"] = username
+                session_details["session_id"] = str(uuid.uuid4()) + str(time.time())
+
                 return (
                     gr.update(visible=False),  # hide login
                     gr.update(visible=True),  # show main
                     gr.update(visible=False),  # hide error
                     True  # set state
                 )
+
             return None, None, gr.update(value="❌ Incorrect username or password", visible=True), False
 
         def clear_chat():
@@ -336,10 +367,8 @@ def create_demo(agent, tools_dict):
             return (
                 gr.update(value=result["Conclusion"], lines=4, interactive=True),
                 gr.update(visible=True),
-                # gr.update(visible=True),
-                )
-        
-        # def records_to_pdf(table, conclusion) -> Path:
+            )
+
         def records_to_pdf(conclusion) -> Path:
             """
             Writes a PDF report under ./reports/  and returns the Path.
@@ -351,28 +380,14 @@ def create_demo(agent, tools_dict):
 
             pdf.cell(0, 10, "Chest-X-ray Report", ln=1, align="C")
             pdf.ln(4)
-
-            # pdf.set_font(family="Helvetica", style="B")
-            # pdf.cell(60, 8, "Disease")
-            # pdf.cell(0, 8, "Information", ln=1)
-            # pdf.set_font(family="Helvetica", style="")
-
-            # for idx, row in table.iterrows():
-            #     pdf.multi_cell(0, 8, f"{row['Disease']}: {row['Info']}")
-
-            # pdf.ln(4)
-            # pdf.set_font(family="Helvetica", style="B")
-            # pdf.cell(0, 8, "Conclusion", ln=1)
             pdf.set_font(family="Helvetica", style="")
             pdf.multi_cell(0, 8, conclusion)
 
             pdf_path = REPORT_DIR / f"report_{uuid.uuid4().hex}.pdf"
             pdf.output(str(pdf_path))
             return pdf_path
-        
-        # def build_pdf_and_preview(table, conclusion):
+
         def build_pdf_and_preview(conclusion):
-            # pdf_path = records_to_pdf(table, conclusion)
             pdf_path = records_to_pdf(conclusion)
 
             iframe_html = (
@@ -433,13 +448,8 @@ def create_demo(agent, tools_dict):
             outputs=[chatbot, image_display, txt],
         ).then(lambda: gr.Textbox(interactive=True), None, [txt])
 
-        # upload_button.upload(handle_file_upload, inputs=upload_button, outputs=[image_display])
-
-        # dicom_upload.upload(handle_file_upload, inputs=dicom_upload, outputs=[image_display])
-
         clear_btn.click(clear_chat, outputs=[chatbot, image_display])
         new_thread_btn.click(new_thread, outputs=[chatbot, image_display, conclusion_tb, approve_btn, download_pdf_btn, pdf_preview])
-        # generate_report_btn.click(generate_report, outputs=[diseases_df, conclusion_tb, approve_btn, reject_btn])
         generate_report_btn.click(generate_report, outputs=[conclusion_tb, approve_btn])
         approve_btn.click(
             build_pdf_and_preview,
@@ -447,7 +457,5 @@ def create_demo(agent, tools_dict):
             inputs=[conclusion_tb],
             outputs=[download_pdf_btn, pdf_preview],
         )
-        # reject_btn.click(show_reject_ui, outputs=[rejection_text, submit_reject_btn, cancel_reject_btn])
-        # cancel_reject_btn.click(hide_reject_ui, outputs=[rejection_text, submit_reject_btn, cancel_reject_btn])
 
     return demo