chatbot.py

import json
import logging
from typing import Dict, List, Optional

import streamlit as st
from snowflake.snowpark.context import get_active_session
from snowflake.snowpark.session import Session

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Available models for selection
AVAILABLE_MODELS = [
    "claude-3-5-sonnet",
    "gemma-7b",
    "jamba-1.5-mini",
    "jamba-1.5-large",
    "jamba-instruct",
    "llama2-70b-chat",
    "llama3-8b",
    "llama3-70b",
    "llama3.1-8b",
    "llama3.1-70b",
    "llama3.1-405b",
    "llama3.2-1b",
    "llama3.2-3b",
    "llama3.3-70b",
    "mistral-large",
    "mistral-large2",
    "mistral-7b",
    "mixtral-8x7b",
    "reka-core",
    "reka-flash",
    "snowflake-arctic",
    "snowflake-llama-3.1-405b",
    "snowflake-llama-3.3-70b",
]


class ChatMessage:
    """
    A class to represent and validate individual chat messages.

    This class ensures that messages conform to expected format and content requirements
    before they are processed by the chat system.
    """

    def __init__(self, role: str, content: str):
        """
        Initialize a new chat message with validation.

        Args:
            role (str): The role of the message sender ('assistant' or 'user')
            content (str): The actual message content

        Raises:
            ValueError: If role is invalid or content is empty/non-string
        """
        # Validate the role is one of the allowed types
        if role not in ["assistant", "user"]:
            raise ValueError(f"Invalid role: {role}")

        # Ensure content is a non-empty string
        if not content or not isinstance(content, str):
            raise ValueError("Content must be a non-empty string")

        self.role = role
        self.content = content

    def to_dict(self) -> Dict[str, str]:
        """
        Convert the message to a dictionary format for API communication.

        Returns:
            Dict[str, str]: A dictionary containing role and content
        """
        return {"role": self.role, "content": self.content}


class SnowflakeChat:
    """
    Handles all Snowflake-specific chat operations including message formatting
    and communication with the Snowflake Cortex API.
    """

    def __init__(self, session: Session):
        """
        Initialize the Snowflake chat handler.

        Args:
            session (Session): An active Snowflake session for database operations
        """
        self.session = session
        # Model will be set dynamically based on user selection
        self.model = None
        # Default parameters for the model
        self.default_params = {
            "temperature": 0.7,  # Controls response randomness (0.0-1.0)
            "max_tokens": 4096,  # Maximum length of generated response
        }

    def set_model(self, model: str):
        """
        Set the model to use for chat completion.

        Args:
            model (str): The name of the model to use
        """
        if model not in AVAILABLE_MODELS:
            raise ValueError(f"Invalid model: {model}")
        self.model = model

    def format_messages(self, messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
        """
        Format a list of messages for the Snowflake Cortex API.

        This method handles message validation and escaping special characters
        to prevent SQL injection and ensure proper JSON formatting.

        Args:
            messages (List[Dict[str, str]]): Raw messages from the chat history

        Returns:
            List[Dict[str, str]]: Formatted messages ready for API submission
        """
        # Skip the first message if it's just the initial greeting
        messages_to_send = messages[1:] if len(messages) > 1 else messages
        formatted_messages = []

        for msg in messages_to_send:
            try:
                # Validate message format using ChatMessage class
                chat_msg = ChatMessage(msg["role"], msg["content"])
                # Escape newlines and quotes for proper JSON formatting
                content = chat_msg.content.replace("\n", "\\n").replace('"', '\\"')
                formatted_messages.append({"role": chat_msg.role, "content": content})
            except ValueError as e:
                # Log formatting errors but continue processing other messages
                logger.error(f"Error formatting message: {e}")
                continue

        return formatted_messages

    def query_cortex(self, messages: List[Dict[str, str]]) -> Optional[str]:
        """
        Send a query to Snowflake Cortex and handle the response.

        This method constructs and executes the SQL query to interact with the
        Snowflake Cortex API, handling any errors that occur during the process.

        Args:
            messages (List[Dict[str, str]]): Formatted messages to send to Cortex

        Returns:
            Optional[str]: The model's response text, or None if an error occurs

        Raises:
            Exception: Various exceptions might be raised during API communication
        """
        if not self.model:
            raise ValueError("Model not set")

        try:
            # Prepare messages for SQL query, escaping single quotes
            messages_json = json.dumps(messages).replace("'", "''")

            # Construct the SQL query for the Cortex API
            query = f"""
            SELECT SNOWFLAKE.CORTEX.COMPLETE(
                '{self.model}',
                parse_json('{messages_json}'),
                parse_json('{json.dumps(self.default_params)}')
            )
            """

            # Execute query and get results
            result = self.session.sql(query).collect()
            response_json = json.loads(result[0][0])

            # Validate response contains expected data
            if not response_json.get("choices"):
                raise ValueError("No choices in response")

            return response_json["choices"][0]["messages"]

        except Exception as e:
            # Log error details and re-raise for higher-level handling
            logger.error(f"Error querying Cortex: {e}")
            raise


class ChatInterface:
    """
    Manages the Streamlit user interface and coordinates between the UI
    and the Snowflake chat functionality.
    """

    def __init__(self):
        """
        Initialize the chat interface, setting up session state and
        establishing Snowflake connection.
        """
        self.initialize_session_state()
        self.snowflake_session = self._get_snowflake_session()
        if self.snowflake_session:
            self.chat = SnowflakeChat(self.snowflake_session)

    @staticmethod
    def initialize_session_state():
        """
        Initialize or reset the Streamlit session state for chat history
        and model selection.
        """
        if "messages" not in st.session_state:
            st.session_state["messages"] = [
                ChatMessage("assistant", "How can I help you?").to_dict()
            ]
        if "selected_model" not in st.session_state:
            st.session_state["selected_model"] = "claude-3-5-sonnet"

    @staticmethod
    def _get_snowflake_session() -> Optional[Session]:
        """
        Establish a connection to Snowflake with error handling.

        Returns:
            Optional[Session]: Active Snowflake session or None if connection fails
        """
        try:
            return get_active_session()
        except Exception as e:
            logger.error(f"Failed to connect to Snowflake: {e}")
            st.error(
                "❌ Failed to connect to Snowflake. Please check your connection settings."
            )
            return None

    def handle_model_change(self):
        """
        Handle model selection changes and update the chat interface accordingly.
        """
        selected_model = st.selectbox(
            "Select Model",
            options=AVAILABLE_MODELS,
            index=AVAILABLE_MODELS.index(st.session_state.selected_model),
            key="model_selector",
        )

        if selected_model != st.session_state.selected_model:
            st.session_state.selected_model = selected_model
            # Clear chat history when model changes
            # st.session_state.messages = [
            #     ChatMessage("assistant", "How can I help you?").to_dict()
            # ]
            # st.re_run()

    def handle_user_input(self):
        """
        Process user input, send it to Snowflake Cortex, and display the response.

        This method handles the main chat interaction loop, including:
        - Capturing user input
        - Updating chat history
        - Processing messages
        - Displaying responses
        - Error handling
        """
        # Set the current model
        self.chat.set_model(st.session_state.selected_model)

        if prompt := st.chat_input("Type your message here..."):
            # Add user's message to chat history
            st.session_state.messages.append(ChatMessage("user", prompt).to_dict())
            st.chat_message("user").write(prompt)

            with st.chat_message("assistant"):
                with st.spinner("💭 Thinking..."):
                    try:
                        # Format messages for API
                        agentic_messages = self.chat.format_messages(
                            st.session_state.messages
                        )

                        # Add instruction for specific answers to the last message
                        if agentic_messages:
                            last_message = agentic_messages[-1]["content"]
                            agentic_messages[-1][
                                "content"
                            ] = f"{last_message} - *Only Answer very specific to this question*"

                        # Get response from Snowflake Cortex
                        response_content = self.chat.query_cortex(agentic_messages)

                        # Update chat history with assistant's response
                        assistant_message = ChatMessage(
                            "assistant", response_content
                        ).to_dict()
                        st.session_state.messages.append(assistant_message)
                        st.write(response_content)

                    except json.JSONDecodeError as e:
                        # Handle JSON parsing errors
                        logger.error(f"JSON parsing error: {e}")
                        st.error("🚫 Error processing the response. Please try again.")
                    except Exception as e:
                        # Handle all other errors
                        logger.error(f"Unexpected error: {e}")
                        st.error(
                            "⚠️ Something went wrong. Please try again or contact support."
                        )
                        # Offer debug information for troubleshooting
                        if st.checkbox("Show debug information"):
                            st.code(str(e))

    def display_chat_history(self):
        """
        Display the entire chat history in the Streamlit interface.
        Each message is displayed with appropriate styling based on its role.
        """
        for msg in st.session_state.messages:
            with st.chat_message(msg["role"]):
                st.write(msg["content"])


def main():
    """
    Main application entry point.

    Sets up the Streamlit page configuration, initializes the chat interface,
    and manages the main application flow.
    """
    # Configure Streamlit page settings
    st.set_page_config(page_title="Snowflake Chat", page_icon="💬", layout="wide")

    # Display application title
    st.title("💬 Snowflake Completion Chatbot")

    # Initialize chat interface
    chat_interface = ChatInterface()

    # Only proceed if Snowflake connection is successful
    if chat_interface.snowflake_session:
        # Add model selector before chat history
        chat_interface.handle_model_change()

        # Display chat interface
        chat_interface.display_chat_history()
        chat_interface.handle_user_input()
    else:
        st.warning("Please ensure you have a valid Snowflake connection to continue.")


if __name__ == "__main__":
    main()