Python/neural_network/chatbot/chatbot.py

import datetime
from typing import List, Dict, Any


class Chatbot:
    """
    A Chatbot class to manage chat conversations using an LLM service and a database to store chat data.

    Methods:
    - start_chat: Starts a new conversation, logs the start time.
    - handle_user_message: Processes user input and stores user message & bot response in DB.
    - end_chat: Ends the conversation and logs the end time.
    - continue_chat: Retains only the last few messages if the conversation exceeds 1000 messages.
    """

    def __init__(self, db: Any, llm_service: Any) -> None:
        """
        Initialize the Chatbot with a database and an LLM service.

        Parameters:
        - db: The database instance used for storing chat data.
        - llm_service: The language model service for generating responses.
        """
        self.db = db
        self.llm_service = llm_service
        self.conversation_history: List[Dict[str, str]] = []
        self.chat_id_pk: int = None

    def start_chat(self) -> None:
        """
        Start a new chat session and insert chat history to the database.
        """
        start_time = datetime.datetime.now()
        is_stream = 1  # Start new conversation
        self.db.insert_chat_history(start_time, is_stream)
        self.chat_id_pk = self.db.get_latest_chat_id()

    def handle_user_message(self, user_input: str) -> str:
        """
        Handle user input and generate a bot response.
        If the user sends '/stop', the conversation is terminated.

        Parameters:
        - user_input: The input provided by the user.

        Returns:
        - bot_response: The response generated by the bot.

        Raises:
        - ValueError: If user input is not a string or if no chat_id is available.

        Doctest:
        >>> class MockDatabase:
        ...     def __init__(self):
        ...         self.data = []
        ...     def insert_chat_data(self, *args, **kwargs):
        ...         pass
        ...     def insert_chat_history(self, *args, **kwargs):
        ...         pass
        ...     def get_latest_chat_id(self):
        ...         return 1
        ...
        >>> class MockLLM:
        ...     def generate_response(self, conversation_history):
        ...         if conversation_history[-1]["content"] == "/stop":
        ...             return "conversation-terminated"
        ...         return "Mock response"
        >>> db_mock = MockDatabase()
        >>> llm_mock = MockLLM()
        >>> bot = Chatbot(db_mock, llm_mock)
        >>> bot.start_chat()
        >>> bot.handle_user_message("/stop")
        'conversation-terminated'
        >>> bot.handle_user_message("Hello!")
        'Mock response'
        """
        if not isinstance(user_input, str):
            raise ValueError("User input must be a string.")

        if self.chat_id_pk is None:
            raise ValueError("Chat has not been started. Call start_chat() first.")

        self.conversation_history.append({"role": "user", "content": user_input})

        if user_input == "/stop":
            bot_response = "conversation-terminated"
            # print(f"Bot: {bot_response}")
            self.end_chat()
            return bot_response
        else:
            bot_response = self.llm_service.generate_response(self.conversation_history)
            # print(f"Bot: {bot_response}")
            self.conversation_history.append(
                {"role": "assistant", "content": bot_response}
            )
            self._store_message_in_db(user_input, bot_response)

        return bot_response

    def _store_message_in_db(self, user_input: str, bot_response: str) -> None:
        """
        Store user input and bot response in the database.

        Parameters:
        - user_input: The message from the user.
        - bot_response: The response generated by the bot.

        Raises:
        - ValueError: If insertion into the database fails.
        """
        try:
            self.db.insert_chat_data(self.chat_id_pk, user_input, bot_response)
        except Exception as e:
            raise ValueError(f"Failed to insert chat data: {e}")

    def end_chat(self) -> None:
        """
        End the chat session and update the chat history in the database.
        """
        current_time = datetime.datetime.now()
        is_stream = 2  # End of conversation
        try:
            user_input = "/stop"
            bot_response = "conversation-terminated"
            self.db.insert_chat_data(self.chat_id_pk, user_input, bot_response)
            self.db.insert_chat_history(current_time, is_stream)
        except Exception as e:
            raise ValueError(f"Failed to update chat history: {e}")

    def continue_chat(self) -> None:
        """
        Retain only the last few entries if the conversation exceeds 1000 messages.
        """
        if len(self.conversation_history) > 1000:
            self.conversation_history = self.conversation_history[-3:]