messages.py

"""
Message classes for the chatbot application.

This module contains the message classes used throughout the app.
By keeping them in a separate module, they remain stable across
Streamlit app reruns, avoiding isinstance comparison issues.
"""
import streamlit as st
from abc import ABC, abstractmethod


class Message(ABC):
    def __init__(self):
        pass

    @abstractmethod
    def to_input_messages(self):
        """Convert this message into a list of dicts suitable for the model API."""
        pass

    @abstractmethod
    def render(self, idx):
        """Render the message in the Streamlit app."""
        pass


class UserMessage(Message):
    def __init__(self, content):
        super().__init__()
        self.content = content

    def to_input_messages(self):
        return [{
            "role": "user",
            "content": self.content
        }]

    def render(self, _):
        with st.chat_message("user"):
            st.markdown(self.content)


class AssistantResponse(Message):
    def __init__(self, messages, request_id):
        super().__init__()
        self.messages = messages
        # Request ID tracked to enable submitting feedback on assistant responses via the feedback endpoint
        self.request_id = request_id

    def to_input_messages(self):
        return self.messages

    def render(self, idx):
        with st.chat_message("assistant"):
            for msg in self.messages:
                render_message(msg)

            if self.request_id is not None:
                render_assistant_message_feedback(idx, self.request_id)


def render_message(msg):
    """Render a single message with healthcare-focused styling."""
    if msg["role"] == "assistant":
        # Render content first if it exists
        if msg.get("content"):
            st.markdown(msg["content"])
        
        # Then render tool calls if they exist with healthcare-focused icons and messaging
        if "tool_calls" in msg and msg["tool_calls"]:
            for call in msg["tool_calls"]:
                fn_name = call["function"]["name"]
                args = call["function"]["arguments"]
                st.markdown(f"""
                <div style="background: #f0f9ff; border-left: 4px solid #0ea5e9; padding: 1rem; margin: 0.5rem 0; border-radius: 0 8px 8px 0;">
                    <strong>⚙️ Healthcare Analysis Tool: {fn_name}</strong>
                    <details>
                        <summary style="cursor: pointer; color: #0c4a6e;">View parameters</summary>
                        <pre style="background: #e0f2fe; padding: 0.5rem; border-radius: 4px; margin-top: 0.5rem; font-size: 0.9rem;"><code>{args}</code></pre>
                    </details>
                </div>
                """, unsafe_allow_html=True)
    elif msg["role"] == "tool":
        st.markdown("""
        <div style="background: #f0fdf4; border-left: 4px solid #22c55e; padding: 1rem; margin: 0.5rem 0; border-radius: 0 8px 8px 0;">
            <strong>📈 Analysis Results:</strong>
        </div>
        """, unsafe_allow_html=True)
        st.code(msg["content"], language="json")


@st.fragment
def render_assistant_message_feedback(i, request_id):
    """Render healthcare-focused feedback UI for assistant messages."""
    from model_serving_utils import submit_feedback
    import os
    
    def save_feedback(index):
        serving_endpoint = os.getenv('SERVING_ENDPOINT')
        if serving_endpoint:
            submit_feedback(
                endpoint=serving_endpoint,
                request_id=request_id,
                rating=st.session_state[f"feedback_{index}"]
            )
    
    # Healthcare-focused feedback prompt
    st.markdown("""
    <div style="background: #f8fafc; border-radius: 8px; padding: 1rem; margin-top: 1rem;">
        <small style="color: #64748b;">Was this healthcare information helpful?</small>
    </div>
    """, unsafe_allow_html=True)
    
    st.feedback("thumbs", key=f"feedback_{i}", on_change=save_feedback, args=[i])