Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py for Gradio App on Hugging Face Spaces
+# app.py for Gradio App on Hugging Face Spaces (with Sandbox Tab)
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -49,8 +49,8 @@ class MedicalSimulatorState:
         self.ordered_tests: Dict[str, str] = {} # e.g., {"cbc": "pending", "xray": "result..."}
         # Add more state variables as needed
 
-# --- Core AI Interaction Function ---
-def get_ai_response(user_input: str, history: List[Tuple[Optional[str], Optional[str]]], patient_profile: Dict, underlying_diagnosis: str) -> str:
+# --- Core AI Interaction Function (Medical Context) ---
+def get_ai_response_medical(user_input: str, history: List[Tuple[Optional[str], Optional[str]]], patient_profile: Dict, underlying_diagnosis: str) -> str:
     if not model or not tokenizer:
         return "Error: AI model is not loaded."
 
@@ -82,6 +82,39 @@ def get_ai_response(user_input: str, history: List[Tuple[Optional[str], Optional
         print(f"Error during AI generation: {e}")
         return f"An error occurred while processing the AI response: {e}"
 
+# --- Core AI Interaction Function (General/Sandbox Context) ---
+def get_ai_response_general(user_input: str, history: List[Tuple[Optional[str], Optional[str]]]) -> str:
+    if not model or not tokenizer:
+        return "Error: AI model is not loaded."
+
+    # Construct a prompt for the AI based on general chat history
+    history_str = "\n".join([f"{'User' if h[0] else 'Assistant'}: {h[0] or h[1]}" for h in history])
+    prompt = f"<|system|>You are a helpful assistant. Answer the user's questions and follow their instructions.<|user|>{history_str}\n{user_input}<|assistant|>"
+
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt")
+        if inputs["input_ids"].shape[1] > 32768: # Check for model max length
+             return "Error: Input prompt is too long for the model."
+
+        # Generate response
+        generate_ids = model.generate(
+            inputs.input_ids,
+            max_new_tokens=512, # Limit generated tokens
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+        # Decode the generated response
+        response_text = tokenizer.decode(generate_ids[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        return response_text.strip()
+
+    except Exception as e:
+        print(f"Error during AI generation: {e}")
+        return f"An error occurred while processing the AI response: {e}"
+
+
 # --- Tool Functions (Modify State) ---
 def start_case(case_type: str, state: MedicalSimulatorState) -> Tuple[MedicalSimulatorState, List[Tuple[Optional[str], Optional[str]]], str, str]:
     # --- Generate Patient Profile (Simplified Example) ---
@@ -148,7 +181,7 @@ def handle_chat(user_input: str, history: List[Tuple[Optional[str], Optional[str
     history.append((user_input, None))
 
     # Get AI response
-    ai_response = get_ai_response(user_input, history[:-1], state.patient_profile, state.underlying_diagnosis) # Pass history without the user's new message yet
+    ai_response = get_ai_response_medical(user_input, history[:-1], state.patient_profile, state.underlying_diagnosis) # Pass history without the user's new message yet
 
     # Add AI response to history
     history[-1] = (user_input, ai_response) # Update the last entry with the AI's response
@@ -164,7 +197,7 @@ def use_tool(tool_name: str, state: MedicalSimulatorState) -> Tuple[MedicalSimul
     state.total_cost += cost
 
     if tool_name == "ask_question":
-        ai_response = get_ai_response("The user asks a general question to gather more history.", state.chat_history, state.patient_profile, state.underlying_diagnosis)
+        ai_response = get_ai_response_medical("The user asks a general question to gather more history.", state.chat_history, state.patient_profile, state.underlying_diagnosis)
         state.chat_history.append(("System", f"[Action: {tool_name}, Cost: ${cost:.2f}]"))
         state.chat_history.append(("AI Patient", ai_response))
 
@@ -214,69 +247,102 @@ def end_case(state: MedicalSimulatorState) -> Tuple[MedicalSimulatorState, List[
     profile_str = "\n".join([f"{k.replace('_', ' ').title()}: {v}" for k, v in (state.patient_profile or {}).items()])
     return state, state.chat_history, f"${state.total_cost:.2f}", profile_str
 
+# --- Sandbox Chat Handler ---
+def handle_sandbox_chat(user_input: str, history: List[Tuple[Optional[str], Optional[str]]]) -> List[Tuple[Optional[str], Optional[str]]]:
+    if not user_input.strip():
+        return history
+
+    # Add user message to history
+    history.append((user_input, None))
+
+    # Get AI response (general context)
+    ai_response = get_ai_response_general(user_input, history[:-1])
+
+    # Add AI response to history
+    history[-1] = (user_input, ai_response)
+
+    return history
+
 
 # --- Gradio Interface ---
 with gr.Blocks(title="Advanced Medical Simulator") as demo:
-    # State component to hold the simulator state across interactions
-    state = gr.State(lambda: MedicalSimulatorState())
-
     gr.Markdown("# Advanced Medical Simulator")
 
-    with gr.Row():
-        with gr.Column(scale=2):
-            # Chat Interface
-            chatbot = gr.Chatbot(label="Patient Interaction", height=400, bubble_full_width=False)
-            with gr.Row():
-                user_input = gr.Textbox(label="Your Action / Question", placeholder="Type your action or question here...", scale=4)
-                submit_btn = gr.Button("Submit", scale=1)
-
-        with gr.Column(scale=1):
-            # Patient Chart / Info
-            patient_chart = gr.Markdown(label="Patient Chart", value="Click 'Start New Case' to begin.")
-            cost_display = gr.Textbox(label="Total Cost", value="$0.00", interactive=False)
-
-    with gr.Row():
-        # Tool Panel
-        with gr.Column():
-            gr.Markdown("### Tools")
-            with gr.Row():
-                ask_btn = gr.Button("Ask Question ($10)")
-                exam_btn = gr.Button("Physical Exam ($25)")
-            with gr.Row():
-                cbc_btn = gr.Button("Order CBC ($50)")
-                xray_btn = gr.Button("Order X-Ray ($150)")
-            with gr.Row():
-                 med_btn = gr.Button("Administer Med ($30)")
-                 end_btn = gr.Button("End Case", variant="stop") # Red button for ending
-
-    with gr.Row():
-        # Case Controls
-        start_case_btn = gr.Button("Start New Case (General)")
-        case_type_dropdown = gr.Dropdown(["General", "Psychiatry", "Pediatric", "Dual Diagnosis"], label="Case Type", value="General")
-
-    # Event Handling
-    start_case_btn.click(
-        fn=start_case,
-        inputs=[case_type_dropdown, state],
-        outputs=[state, chatbot, cost_display, patient_chart]
-    )
-
-    submit_btn.click(
-        fn=handle_chat,
-        inputs=[user_input, chatbot, state],
-        outputs=[chatbot, cost_display]
-    ).then(
-         fn=lambda: "", # Clear the input textbox after submission
-         inputs=[],
-         outputs=[user_input]
-    )
-
-    ask_btn.click(fn=lambda s: use_tool("ask_question", s), inputs=[state], outputs=[state, chatbot, cost_display])
-    exam_btn.click(fn=lambda s: use_tool("physical_exam", s), inputs=[state], outputs=[state, chatbot, cost_display])
-    cbc_btn.click(fn=lambda s: use_tool("order_cbc", s), inputs=[state], outputs=[state, chatbot, cost_display])
-    xray_btn.click(fn=lambda s: use_tool("order_xray", s), inputs=[state], outputs=[state, chatbot, cost_display])
-    med_btn.click(fn=lambda s: use_tool("administer_med", s), inputs=[state], outputs=[state, chatbot, cost_display])
-    end_btn.click(fn=end_case, inputs=[state], outputs=[state, chatbot, cost_display, patient_chart])
+    with gr.Tab("Medical Simulation"):
+        # State component to hold the simulator state across interactions
+        state = gr.State(lambda: MedicalSimulatorState())
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                # Chat Interface
+                chatbot = gr.Chatbot(label="Patient Interaction", height=400, bubble_full_width=False)
+                with gr.Row():
+                    user_input = gr.Textbox(label="Your Action / Question", placeholder="Type your action or question here...", scale=4)
+                    submit_btn = gr.Button("Submit", scale=1)
+
+            with gr.Column(scale=1):
+                # Patient Chart / Info
+                patient_chart = gr.Markdown(label="Patient Chart", value="Click 'Start New Case' to begin.")
+                cost_display = gr.Textbox(label="Total Cost", value="$0.00", interactive=False)
+
+        with gr.Row():
+            # Tool Panel
+            with gr.Column():
+                gr.Markdown("### Tools")
+                with gr.Row():
+                    ask_btn = gr.Button("Ask Question ($10)")
+                    exam_btn = gr.Button("Physical Exam ($25)")
+                with gr.Row():
+                    cbc_btn = gr.Button("Order CBC ($50)")
+                    xray_btn = gr.Button("Order X-Ray ($150)")
+                with gr.Row():
+                     med_btn = gr.Button("Administer Med ($30)")
+                     end_btn = gr.Button("End Case", variant="stop") # Red button for ending
+
+        with gr.Row():
+            # Case Controls
+            start_case_btn = gr.Button("Start New Case (General)")
+            case_type_dropdown = gr.Dropdown(["General", "Psychiatry", "Pediatric", "Dual Diagnosis"], label="Case Type", value="General")
+
+        # Event Handling for Medical Simulation Tab
+        start_case_btn.click(
+            fn=start_case,
+            inputs=[case_type_dropdown, state],
+            outputs=[state, chatbot, cost_display, patient_chart]
+        )
+
+        submit_btn.click(
+            fn=handle_chat,
+            inputs=[user_input, chatbot, state],
+            outputs=[chatbot, cost_display]
+        ).then(
+             fn=lambda: "", # Clear the input textbox after submission
+             inputs=[],
+             outputs=[user_input]
+        )
+
+        ask_btn.click(fn=lambda s: use_tool("ask_question", s), inputs=[state], outputs=[state, chatbot, cost_display])
+        exam_btn.click(fn=lambda s: use_tool("physical_exam", s), inputs=[state], outputs=[state, chatbot, cost_display])
+        cbc_btn.click(fn=lambda s: use_tool("order_cbc", s), inputs=[state], outputs=[state, chatbot, cost_display])
+        xray_btn.click(fn=lambda s: use_tool("order_xray", s), inputs=[state], outputs=[state, chatbot, cost_display])
+        med_btn.click(fn=lambda s: use_tool("administer_med", s), inputs=[state], outputs=[state, chatbot, cost_display])
+        end_btn.click(fn=end_case, inputs=[state], outputs=[state, chatbot, cost_display, patient_chart])
+
+    with gr.Tab("AI Sandbox"):
+        sandbox_chatbot = gr.Chatbot(label="General AI Chat", height=500, bubble_full_width=False)
+        with gr.Row():
+            sandbox_input = gr.Textbox(label="Message", placeholder="Ask anything...", scale=4)
+            sandbox_submit = gr.Button("Send", scale=1)
+
+        sandbox_submit.click(
+            fn=handle_sandbox_chat,
+            inputs=[sandbox_input, sandbox_chatbot],
+            outputs=[sandbox_chatbot]
+        ).then(
+             fn=lambda: "", # Clear the input textbox after submission
+             inputs=[],
+             outputs=[sandbox_input]
+        )
 
 # Launch the app
 # For Hugging Face Spaces, Gradio handles the launch.