testing conversational abilities

app.py

@@ -102,7 +102,7 @@ class RoleAgent:
                 thinking = block.split("ANSWER:", 1)[0].strip()
                 answer = block.split("ANSWER:", 1)[1].strip()
         print(
-            "[DEBUG] thinking/answer split:",
+            "[THINKING ANSWER SPLIT] thinking/answer split:",
             response,
             "→",
             "[THINKING] thinking:",
@@ -157,91 +157,142 @@ diagnosis = ""
 
 
 # === Gradio Inference ===
-def simulate_interaction(user_input, iterations=1):
-    history = [f"Doctor: What brings you in today?", f"Patient: {user_input}"]
-    summary, diagnosis = "", ""
-    treatment_out = None
+def simulate_interaction(user_input, conversation_history=None):
+    """Single turn interaction - no iterations, uses accumulated history"""
+    if conversation_history is None:
+        history = [f"Doctor: What brings you in today?", f"Patient: {user_input}"]
+    else:
+        history = conversation_history.copy()
+        history.append(f"Patient: {user_input}")
 
-    for i in range(iterations):
-        # Summarize
-        sum_in = "\n".join(history) + f"\nPrevious Vignette: {summary}"
-        sum_out = summarizer.act(sum_in)
-        summary = sum_out["output"]
+    # Summarize the full conversation history
+    sum_in = "\n".join(history)
+    sum_out = summarizer.act(sum_in)
+    summary = sum_out["output"]
 
-        # Diagnose
-        diag_out = diagnoser.act(summary)
-        diagnosis = diag_out["output"]
+    # Diagnose based on summary
+    diag_out = diagnoser.act(summary)
+    diagnosis = diag_out["output"]
 
-        # Question
-        q_in = f"Vignette: {summary}\n" f"Current Estimated Diagnosis:\n" f"{diagnosis}"
-        q_out = questioner.act(q_in)
-        history.append(f"Doctor: {q_out['output']}")
+    # Generate next question based on current understanding
+    q_in = f"Vignette: {summary}\nCurrent Estimated Diagnosis: {diagnosis}"
+    q_out = questioner.act(q_in)
 
-        # Treatment
-        treatment_out = treatment_agent.act(f"Diagnosis: {diagnosis}\n")
+    # Add doctor's response to history
+    history.append(f"Doctor: {q_out['output']}")
+
+    # Generate treatment plan (but don't end conversation)
+    treatment_out = treatment_agent.act(f"Diagnosis: {diagnosis}\nVignette: {summary}")
 
     return {
         "summary": sum_out,
         "diagnosis": diag_out,
         "question": q_out,
         "treatment": treatment_out,
-        "conversation": "\n".join(history),
+        "conversation": history,  # Return full history list
     }
 
 
 # === Gradio UI ===
 def ui_fn(user_input):
+    """Non-stateful version for testing"""
     res = simulate_interaction(user_input)
     return f"""📋 Vignette Summary:
 💭 THINKING: {res['summary']['thinking']}
-ANSWER: {res['summary']['output']}
+📝 SUMMARY: {res['summary']['output']}
 
 🩺 Diagnosis:
 💭 THINKING: {res['diagnosis']['thinking']}
-ANSWER: {res['diagnosis']['output']}
-T
+🔍 DIAGNOSIS: {res['diagnosis']['output']}
+
 ❓ Follow-up Question:
 💭 THINKING: {res['question']['thinking']}
-ANSWER: {res['question']['output']}
+👨‍⚕️ DOCTOR: {res['question']['output']}
 
 💊 Treatment Plan:
-{res['treatment']['output']}
+💭 THINKING: {res['treatment']['thinking']}
+📋 TREATMENT: {res['treatment']['output']}
 
-💬 Conversation:
-{res['conversation']}
+💬 Full Conversation:
+{chr(10).join(res['conversation'])}
 """
 
 
 # === Stateful Gradio UI ===
 def stateful_ui_fn(user_input, history):
-    # Initialize or retrieve history list
-    history = history or []
-    # Append the new patient utterance
-    history.append(f"Patient: {user_input}")
-    # Run one round of simulation
-    res = simulate_interaction(user_input)
-    # Extract last doctor line from the fresh conversation
-    last_line = res["question"]["output"]
-    # Append the doctor's new line
-    history.append(last_line)
-    # Build the displayed conversation
-    convo = "\n".join(history)
-    # Return both the display text and updated history
-    return convo, history
+    """Proper stateful conversation handler"""
+    # Initialize history if first interaction
+    if history is None:
+        history = []
+
+    # Run one turn of interaction with accumulated history
+    res = simulate_interaction(user_input, history)
+
+    # Get the updated conversation history
+    updated_history = res["conversation"]
+
+    # Format the display output
+    display_output = f"""💬 Conversation:
+{chr(10).join(updated_history)}
+
+📋 Current Assessment:
+🔍 Diagnosis: {res['diagnosis']['output']}
+💊 Treatment Plan: {res['treatment']['output']}
+"""
+
+    # Return display text and updated history for next turn
+    return display_output, updated_history
+
 
+def chat_interface(user_input, history):
+    """Alternative chat-style interface"""
+    if history is None:
+        history = []
 
-demo = gr.Interface(
+    # Run interaction
+    res = simulate_interaction(user_input, history)
+    updated_history = res["conversation"]
+
+    # Return just the doctor's latest response and updated history
+    doctor_response = res["question"]["output"]
+
+    return doctor_response, updated_history
+
+
+# Create two different interfaces
+demo_stateful = gr.Interface(
     fn=stateful_ui_fn,
     inputs=[
-        gr.Textbox(label="Patient Response"),
+        gr.Textbox(
+            label="Patient Response",
+            placeholder="Describe your symptoms or answer the doctor's question...",
+        ),
         gr.State(),  # holds the conversation history
     ],
     outputs=[
-        gr.Textbox(label="Doctor Simulation Output"),
+        gr.Textbox(label="Medical Consultation", lines=15),
         gr.State(),  # returns the updated history
     ],
-    title="🧠 AI Doctor Multi-Agent Reasoning",
+    title="🧠 AI Doctor - Full Medical Consultation",
+    description="Have a conversation with an AI doctor. Each response builds on the previous conversation.",
+)
+
+demo_chat = gr.Interface(
+    fn=chat_interface,
+    inputs=[
+        gr.Textbox(label="Your Response", placeholder="Tell me about your symptoms..."),
+        gr.State(),
+    ],
+    outputs=[
+        gr.Textbox(label="Doctor", lines=5),
+        gr.State(),
+    ],
+    title="🩺 AI Doctor Chat",
+    description="Simple chat interface with the AI doctor.",
 )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    # Launch the stateful version by default
+    demo_stateful.launch(share=True)
+    # Uncomment the line below to use the chat version instead:
+    # demo_chat.launch(share=True)