updated chatbot.py

neural_network/chatbot/README.md

@@ -0,0 +1,55 @@
+# Chatbot with Chat history stored in Database
+
+This project is a simple chatbot application built using Python, integrating a database for chat history storage and a language model service to generate responses. The chatbot can handle user messages, manage chat history, and terminate conversations upon receiving a `/stop` command.
+
+## Features
+- **Conversation Handling**: The bot processes user inputs and generates responses using a language model service.
+- **Database Integration**: Stores chat data (user messages and bot responses) and maintains chat history.
+- **Session Management**: Supports starting and terminating chat sessions, including proper logging of start and end times.
+- **Message Truncation**: Limits conversation history to the last few messages if the conversation exceeds a large number of entries.
+
+## Components
+- **`Chatbot` Class**: Core logic for handling user messages and managing the chat lifecycle.
+- **`Database` (Mocked in tests)**: Handles chat data storage (methods for inserting and retrieving data).
+- **`LLM Service` (Mocked in tests)**: Generates responses to user input based on conversation history.
+
+## Installation
+1. Clone the repository:
+2. Install the necessary dependencies
+   ```bash
+   pip3 install requirements.txt
+   ```
+4. Run the bot or test it using `doctest`:
+    ```bash
+    python3 -m doctest -v chatbot.py
+    ```
+
+## Usage
+1. **Create Database**: Create a databse named `ChatDB` in Mysql
+2. **Create .env**:
+```
+  # Together API key
+  TOGETHER_API_KEY="YOUR_API_KEY"
+
+  # Groq API key
+  GROQ_API_KEY = "YOUR_API_KEY"
+
+  # MySQL connectionDB (if you're running locally)
+  DB_USER = "<DB_USER_NAME>"
+  DB_PASSWORD = "<DB_USER_NAME>"
+  DB_HOST = "127.0.0.1"
+  DB_NAME = "ChatDB"
+  PORT = "3306"
+ ```
+7. **Handling Messages**: run below command to start the chat in console, you can login to your Database to check the chat history
+```python
+python3 main.py
+```
+10. **Ending the Chat**: When the user sends `/stop`, the chat will terminate and log the end of the conversation with the message 'conversation-terminated'
+
+## Testing
+The code includes basic `doctests` to verify the chatbot's functionality using mock services for the database and language model:
+- Run the tests:
+    ```bash
+    python3 -m doctest -v chatbot.py
+    ```

neural_network/chatbot/chatbot.py

@@ -0,0 +1,135 @@
+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:]

neural_network/chatbot/db.py

@@ -0,0 +1,199 @@
+import os
+from dotenv import load_dotenv
+import mysql.connector
+from mysql.connector import MySQLConnection
+
+load_dotenv()
+
+
+class Database:
+    """
+    A class to manage the connection to the MySQL database using configuration from environment variables.
+
+    Attributes:
+    -----------
+    config : dict
+        The database connection parameters like user, password, host, and database name.
+    """
+
+    def __init__(self) -> None:
+        self.config = {
+            "user": os.environ.get("DB_USER"),
+            "password": os.environ.get("DB_PASSWORD"),
+            "host": os.environ.get("DB_HOST"),
+            "database": os.environ.get("DB_NAME"),
+        }
+
+    def connect(self) -> MySQLConnection:
+        """
+        Establish a connection to the MySQL database.
+
+        Returns:
+        --------
+        MySQLConnection
+            A connection object for interacting with the MySQL database.
+
+        Raises:
+        -------
+        mysql.connector.Error
+            If the connection to the database fails.
+        """
+        return mysql.connector.connect(**self.config)
+
+
+class ChatDatabase:
+    """
+    A class to manage chat-related database operations, such as creating tables,
+    inserting chat history, and retrieving chat data.
+
+    Attributes:
+    -----------
+    db : Database
+        An instance of the `Database` class for establishing connections to the MySQL database.
+    """
+
+    def __init__(self, db: Database) -> None:
+        self.db = db
+
+    def create_tables(self) -> None:
+        """
+        Create the necessary tables for chat history and chat data in the database.
+        If the tables already exist, they will not be created again.
+
+        Raises:
+        -------
+        mysql.connector.Error
+            If there is any error executing the SQL statements.
+        """
+        conn = self.db.connect()
+        cursor = conn.cursor()
+
+        cursor.execute(
+            """
+        CREATE TABLE IF NOT EXISTS ChatDB.Chat_history (
+            chat_id INT AUTO_INCREMENT PRIMARY KEY,
+            start_time DATETIME,
+            is_stream INT
+        )
+        """
+        )
+
+        cursor.execute(
+            """
+        CREATE TABLE IF NOT EXISTS ChatDB.Chat_data (
+            id INT AUTO_INCREMENT PRIMARY KEY,
+            chat_id INT,
+            user TEXT,
+            assistant TEXT,
+            FOREIGN KEY (chat_id) REFERENCES ChatDB.Chat_history(chat_id)
+        )
+        """
+        )
+
+        cursor.execute("DROP TRIGGER IF EXISTS update_is_stream")
+
+        cursor.execute(
+            """
+        CREATE TRIGGER update_is_stream
+        AFTER UPDATE ON ChatDB.Chat_history
+        FOR EACH ROW
+        BEGIN
+            UPDATE ChatDB.Chat_data
+            SET is_stream = NEW.is_stream
+            WHERE chat_id = NEW.chat_id;
+        END;
+        """
+        )
+
+        conn.commit()
+        cursor.close()
+        conn.close()
+
+    def insert_chat_history(self, start_time: str, is_stream: int) -> None:
+        """
+        Insert a new chat history record into the database.
+
+        Parameters:
+        -----------
+        start_time : str
+            The starting time of the chat session.
+        is_stream : int
+            An integer indicating whether the chat is in progress (1) or ended (2).
+
+        Raises:
+        -------
+        mysql.connector.Error
+            If there is any error executing the SQL statements.
+        """
+        conn = self.db.connect()
+        cursor = conn.cursor()
+        cursor.execute(
+            """
+            INSERT INTO ChatDB.Chat_history (start_time, is_stream)
+            VALUES (%s, %s)
+        """,
+            (start_time, is_stream),
+        )
+        conn.commit()
+        cursor.close()
+        conn.close()
+
+    def get_latest_chat_id(self) -> int:
+        """
+        Retrieve the chat ID of the most recent chat session from the database.
+
+        Returns:
+        --------
+        int
+            The ID of the latest chat session.
+
+        Raises:
+        -------
+        mysql.connector.Error
+            If there is any error executing the SQL statements.
+        """
+        conn = self.db.connect()
+        cursor = conn.cursor()
+        cursor.execute(
+            """
+            SELECT chat_id FROM ChatDB.Chat_history WHERE
+            chat_id=(SELECT MAX(chat_id) FROM ChatDB.Chat_history)
+        """
+        )
+        chat_id_pk = cursor.fetchone()[0]
+        cursor.close()
+        conn.close()
+        return chat_id_pk
+
+    def insert_chat_data(
+        self, chat_id: int, user_message: str, assistant_message: str
+    ) -> None:
+        """
+        Insert a new chat data record into the database.
+
+        Parameters:
+        -----------
+        chat_id : int
+            The ID of the chat session to which this data belongs.
+        user_message : str
+            The message provided by the user in the chat session.
+        assistant_message : str
+            The response from the assistant in the chat session.
+
+        Raises:
+        -------
+        mysql.connector.Error
+            If there is any error executing the SQL statements.
+        """
+        conn = self.db.connect()
+        cursor = conn.cursor()
+        cursor.execute(
+            """
+            INSERT INTO ChatDB.Chat_data (chat_id, user, assistant)
+            VALUES (%s, %s, %s)
+        """,
+            (chat_id, user_message, assistant_message),
+        )
+        conn.commit()
+        cursor.close()
+        conn.close()