mathdatech/Python
1
0
Fork 0
mirror of https://github.com/TheAlgorithms/Python.git synced 2025-03-19 21:19:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse Source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2024-10-20 07:49:51 +00:00
parent 322434dc16
commit 5d91b30f87
1 changed files with 37 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
neural_network/chatbot/chat_db.py
View File

@ -22,6 +22,7 @@ db_config = {
"database": os.environ.get("DB_NAME"), "database": os.environ.get("DB_NAME"),
} }
class LLMService: class LLMService:
def __init__(self, api_service: str): def __init__(self, api_service: str):
self.api_service = api_service self.api_service = api_service
@ -62,7 +63,7 @@ class LLMService:
stop=["<|eot_id|>", "<|eom_id|>"], stop=["<|eot_id|>", "<|eom_id|>"],
stream=False, stream=False,
) )
return response.choices[0].message.content return response.choices[0].message.content
@ -176,7 +177,7 @@ class ChatDB:
cursor = conn.cursor() cursor = conn.cursor()
cursor.execute( cursor.execute(
""" """
SELECT chat_id FROM ChatDB.Chat_history SELECT chat_id FROM ChatDB.Chat_history
ORDER BY chat_id DESC LIMIT 1 ORDER BY chat_id DESC LIMIT 1
""" """
) )
@ -190,10 +191,12 @@ class ChatDB:
conn.close() conn.close()
@staticmethod @staticmethod
def insert_chat_data(chat_id: int, user_message: str, assistant_message: str) -> None: def insert_chat_data(
chat_id: int, user_message: str, assistant_message: str
) -> None:
""" """
Insert a new row into the ChatDB.Chat_data table. Insert a new row into the ChatDB.Chat_data table.
Example: Example:
>>> ChatDB.insert_chat_data(1, 'Hello', 'Hi there!') >>> ChatDB.insert_chat_data(1, 'Hello', 'Hi there!')
Chat data inserted successfully. Chat data inserted successfully.
@ -228,7 +231,7 @@ class Chatbot:
""" """
Start a chatbot session, allowing the user to interact with the LLM. Start a chatbot session, allowing the user to interact with the LLM.
Saves conversation history in the database and ends the session on "/stop" command. Saves conversation history in the database and ends the session on "/stop" command.
Example: Example:
>>> chatbot = Chatbot(api_service="Groq") >>> chatbot = Chatbot(api_service="Groq")
>>> chatbot.chat_session() # This will be mocked in the tests >>> chatbot.chat_session() # This will be mocked in the tests
@ -243,11 +246,17 @@ class Chatbot:
if self.chat_id_pk is None: if self.chat_id_pk is None:
if user_input.lower() == "/stop": if user_input.lower() == "/stop":
break break
bot_response = self.llm_service.generate_response(self.conversation_history) bot_response = self.llm_service.generate_response(
self.conversation_history.append({"role": "assistant", "content": bot_response}) self.conversation_history
)
self.conversation_history.append(
{"role": "assistant", "content": bot_response}
)
is_stream = 1 # New conversation is_stream = 1 # New conversation
self.chat_id_pk = ChatDB.insert_chat_history(self.start_time, is_stream) # Return the chat_id self.chat_id_pk = ChatDB.insert_chat_history(
self.start_time, is_stream
) # Return the chat_id
if self.chat_id_pk: if self.chat_id_pk:
ChatDB.insert_chat_data(self.chat_id_pk, user_input, bot_response) ChatDB.insert_chat_data(self.chat_id_pk, user_input, bot_response)
else: else:
@ -257,8 +266,12 @@ class Chatbot:
ChatDB.insert_chat_history(current_time, is_stream) ChatDB.insert_chat_history(current_time, is_stream)
break break
bot_response = self.llm_service.generate_response(self.conversation_history) bot_response = self.llm_service.generate_response(
self.conversation_history.append({"role": "assistant", "content": bot_response}) self.conversation_history
)
self.conversation_history.append(
{"role": "assistant", "content": bot_response}
)
is_stream = 0 # Continuation of conversation is_stream = 0 # Continuation of conversation
current_time = datetime.datetime.now(datetime.timezone.utc) current_time = datetime.datetime.now(datetime.timezone.utc)
@ -268,23 +281,30 @@ class Chatbot:
if len(self.conversation_history) > 1000: if len(self.conversation_history) > 1000:
self.conversation_history = self.conversation_history[-3:] self.conversation_history = self.conversation_history[-3:]
# Test cases for Chatbot # Test cases for Chatbot
class TestChatbot(unittest.TestCase): class TestChatbot(unittest.TestCase):
@patch("builtins.input", side_effect=["Hello", "/stop"])
@patch('builtins.input', side_effect=["Hello", "/stop"]) @patch("sys.stdout", new_callable=StringIO)
@patch('sys.stdout', new_callable=StringIO)
def test_chat_session(self, mock_stdout, mock_input): def test_chat_session(self, mock_stdout, mock_input):
""" """
Test the chat_session method for expected welcome message. Test the chat_session method for expected welcome message.
""" """
chatbot = Chatbot(api_service="Groq") chatbot = Chatbot(api_service="Groq")
chatbot.chat_session() chatbot.chat_session()
# Check for the welcome message in the output # Check for the welcome message in the output
output = mock_stdout.getvalue().strip().splitlines() output = mock_stdout.getvalue().strip().splitlines()
self.assertIn("Welcome to the chatbot! Type '/stop' to end the conversation.", output) self.assertIn(
self.assertTrue(any("Chat history inserted successfully." in line for line in output)) "Welcome to the chatbot! Type '/stop' to end the conversation.", output
self.assertTrue(any("Chat data inserted successfully." in line for line in output)) )
self.assertTrue(
any("Chat history inserted successfully." in line for line in output)
)
self.assertTrue(
any("Chat data inserted successfully." in line for line in output)
)
if __name__ == "__main__": if __name__ == "__main__":
# #