Delete neural_network/chatbot/db.py

This commit is contained in:
Pritam Das 2024-10-20 11:42:34 +05:30 committed by GitHub
parent 922a2306b5
commit a1d4cd9839
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,199 +0,0 @@
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()