app.py

"""
Medical Chatbot using Ollama and Gradio
This template provides a complete chatbot with sections to customize.
"""

import ollama
import gradio as gr

# 1: Customize system prompt for either healthcare professional or patient based on a radio button selection

import gradio as gr
current_mode = "patient"  # default
grExamples = None

def switch_prompt(mode):
    """Switch between patient and doctor mode"""
    global current_mode, SYSTEM_PROMPT
    current_mode = mode
    
    # patient mode
    if mode == "patient":
        SYSTEM_PROMPT = """You are a medical AI assistant specialized in medication questions and management for a patient.

            Your role is to help the patient keep track of their medication schedule and provide information about their medications.

            When answering:
            - Tell them they can let you know about any changes in medications at any time.
            - Let them know about any upcoming medication times.
            - Ask what medications they have taken today so far.
            - Be accurate and evidence-based.
            - Use clear, simple language.

            Remember: You provide educational information, not personal medical advice."""
        
        # Update the example prompts.
    else:  # doctor mode
        SYSTEM_PROMPT = """You are a medical AI assistant specialized in providing information to healthcare professionals about a patient's medications and management.

        Your role is to inform the healthcare professional so they understand the patient's medication list, recent medication schedule, and any concerns the patient has asked about their medications.

        When answering:
        - Present the list of medications the patient is taking.
        - Provide a summary of specific questions the patient has asked.
        - Detail the schedule of the patient's medications as best as you know, and explain what you don't know.

        Remember: You provide educational information, not personal medical advice."""
    
    
    return f"Switched to {mode} mode"

# =============================================
# CONFIGURATION
# =============================================

# Model selection
MODEL = 'llama3.2:3b' # x your Ollama model name here after downloading it

# 2: Set your parameters (tune these for your domain!) try different values and see what works best
PARAMS = {
    'temperature': 0.1,     # Adjust (0.1 = focused, 0.9 = creative)
    'top_p': 0.9,           # Adjust (0.85-0.95 recommended)
    'num_predict': 300,     # Adjust max length
    'repeat_penalty': 1.0   # Adjust if needed
}

# =============================================
# MEMORY & STATE
# =============================================

# Conversation history (for Feature: Memory)
conversation_history = []

# =============================================
# CORE FUNCTIONS
# =============================================

def chat(user_message):
    """
    Main chat function with conversation memory
    
    3: Add your custom logic here
    Ideas:
    - Pre-process user input
    - Add domain-specific checks
    - Format output in special ways
    - Add citations or references
    """
    
    # Validate input
    if not user_message.strip():
        return "Please ask a question!"
    
    # 4: Add any pre-processing here
    
    # If emergency is mentioned, instruct user to seek immediate help
    if "emergency" in user_message.lower() or "immediate" in user_message.lower():
        return "Please seek immediate help from a healthcare professional."
    
    # Add user message to history
    conversation_history.append({
        'role': 'user',
        'content': user_message
    })
    
    # Get AI response
    try:
        response = ollama.chat(
            model=MODEL,
            messages=[
                {'role': 'system', 'content': SYSTEM_PROMPT}
            ] + conversation_history,
            options=PARAMS
        )
        
        # Extract answer
        answer = response['message']['content']
        
        # Add to history
        conversation_history.append({
            'role': 'assistant',
            'content': answer
        })
        
        # 5: Add any post-processing here
        
        # Add disclaimer
        disclaimer = "\n\n" + "─"*50 + "\n**Disclaimer:** This is educational information only. Always consult qualified healthcare professionals for medical advice."
        
        return answer + disclaimer
        
    except Exception as e:
        return f"Error: {str(e)}\n\nMake sure Ollama is running and the model is downloaded."

def clear_conversation():
    """Clear conversation history"""
    global conversation_history
    conversation_history = []
    return "", "Conversation cleared! Start fresh."

# 6: Add your custom feature functions here

def switch_examples(mode):
    """Switch example prompts based on mode"""
    new_dataset = get_examples(mode)
    print ("Switching examples to:", new_dataset.samples)
    return new_dataset

def export_conversation():
    """Export chat history to text file"""
    if not conversation_history:
        return "No conversation to export!"
    
    text = ""
    for msg in conversation_history:
        role = msg['role'].upper()
        content = msg['content']
        text += f"{role}:\n{content}\n\n" + "─"*50 + "\n\n"

    # Save to file
    with open("conversation_export.txt", "w") as f:
        f.write(text)

    return "Exported to conversation_export.txt"

# Create a dictionary to hold the raw lists
EXAMPLE_SAMPLES = {
    "patient": [
        ["How should I manage my blood pressure medication schedule?"],
        ["Why did my doctor change my dosage?"],
        ["Is it normal to feel dizzy after starting this?"],
        ["What are the side effects of my current medications?"],
        ["Which medications should I take with food?"],
        ["What are the interactions between my medications?"],
    ],
    "doctor": [
        ["What is the patient's current medication list?"],
        ["Can you provide a summary of the patient's recent medication schedule?"],
        ["What are the patient's concerns about their medications?"],
    ]
}

def get_examples(mode):
    # Return the correct examples to the gradio component.
    return gr.Dataset(samples=EXAMPLE_SAMPLES[mode])

# =============================================
# GRADIO INTERFACE
# =============================================

# 7: Customize your interface design
with gr.Blocks(theme=gr.themes.Soft(), title="Patient/Doctor Medication Management Chatbot") as app:
    
    # Header
    gr.Markdown(f"""
    # Patient/Doctor Medication Management Chatbot
    ### Patients can ask questions about their medications and schedule. Healthcare professionals can get summaries of the patient's medication experience.
    
    **Domain:** Medication Management 
    **Powered by:** {MODEL} via Ollama
    """)
    
    # Main interface
    with gr.Row():
        with gr.Column(scale=2):
            # Input area
            question = gr.Textbox(
                label="💬 Your Question",
                placeholder="Ask anything about your medications...",
                lines=3
            )
            
            # Buttons
            with gr.Row():
                submit_btn = gr.Button("🚀 Ask", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ Clear", size="lg")
            
            # 8: Add your custom controls here
            # Add button:
            export_btn = gr.Button("Export Conversation")
            export_status = gr.Textbox(label="Export Status")
            export_btn.click(export_conversation, outputs=export_status)

            with gr.Row():
                with gr.Column(scale=2):
                    # A selector for patient or doctor mode.
                    mode_selector = gr.Radio(
                        choices=["patient", "doctor"],
                        value="patient",
                        label="Mode"
                )
            
        with gr.Column(scale=3):
            # Response area
            response = gr.Textbox(
                label="🤖 AI Response",
                lines=15,
                interactive=False
            )

    grExamples = gr.Examples(
        examples=EXAMPLE_SAMPLES[current_mode],
        inputs=question,
        label="💡 Try these examples:",
        cache_examples=False # Required for dynamic swapping
    )

    mode_selector.change(switch_prompt, inputs=mode_selector)
    mode_selector.change(switch_examples, inputs=mode_selector, outputs=grExamples.dataset)
    
    # 10: Add an "About" section
    with gr.Accordion("ℹ️ About This Chatbot", open=False):
        gr.Markdown("""
        ### About
        This chatbot is designed with two modes, patient and doctor, to assist patients and healthcare professionals with medication management. 
        In patient mode, it provides a patient information about medication schedules, side effects, interactions, and more based on their prompts.
        In doctor mode, it provides the doctor or healthcare professional a summary of what the patient has asked about regarding their medications.
        The chatbot will remember the conversation history to provide context-aware responses.
   
        ### Features
        - Feature 1: Answer patient questions about medications and remembers history.
        - Feature 2: Provide doctors with summaries of patient inquiries to help with medication management.
        - Feature 3: Suggests prompts depending on the mode selected (different prompts for doctors and patients).
        - Feature 4: Export conversation history for record-keeping or further analysis.
        
        ### Technical Details
        - **Model:** {model}
        - **Framework:** Ollama + Gradio
        - **Temperature:** {temp}
        - **Max Length:** {max_len} tokens
        
        ### Important Notes
        ⚠️ This chatbot is for educational purposes only.  
        ⚠️ It is NOT a substitute for professional medical advice.  
        ⚠️ Always consult qualified healthcare providers for medical concerns.
        
        ### Creator
        **Developed by:** Donnie Goodson
        **Course:** Medical AI  
        **Date:** April 18, 2026
        """.format(
            model=MODEL,
            temp=PARAMS['temperature'],
            max_len=PARAMS['num_predict']
        ))
    
    # Connect buttons to functions
    submit_btn.click(
        fn=chat,
        inputs=question,
        outputs=response
    )
    
    clear_btn.click(
        fn=clear_conversation,
        inputs=None,
        outputs=[question, response]
    )

    # =============================================
# LAUNCH
# =============================================

if __name__ == "__main__":
    print("\n" + "="*60)
    print("LAUNCHING MEDICAL CHATBOT")
    print("="*60)
    print(f"\nModel: {MODEL}")
    print(f"Domain: Medication Management")
    print(f"Parameters: {PARAMS}")
    print("\nStarting server... Open the URL below in your browser.\n")
    
    app.launch(
        share=True,  # Set to True to get public URL and False for local only
          )