ChatBots in a Jupyter Notebook
Let's see how to use the ChatBot class to enable you to chat with Mistral inside a Jupyter notebook.
from llamabot import ChatBot
code_tester = ChatBot(
"""
You are a Python quality assurance developer who delivers high quality unit tests for code.
You write tests using PyTest and not the built-in unittest library.
Write the tests using test functions and not using classes and class methods
Here is the code to write tests against:
""",
session_name="code-tested",
model_name="mistral/mistral-medium",
stream_target="stdout",
)
code_tester(
'''
class ChatBot:
"""Chat Bot that is primed with a system prompt, accepts a human message.
Automatic chat memory management happens.
h/t Andrew Giessel/GPT4 for the idea.
"""
def __init__(self, system_prompt, temperature=0.0, model_name="gpt-4"):
"""Initialize the ChatBot.
:param system_prompt: The system prompt to use.
:param temperature: The model temperature to use.
See https://platform.openai.com/docs/api-reference/completions/create#completions/create-temperature
for more information.
:param model_name: The name of the OpenAI model to use.
"""
self.model = ChatOpenAI(
model_name=model_name,
temperature=temperature,
streaming=True,
verbose=True,
callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),
)
self.chat_history = [
SystemMessage(content="Always return Markdown-compatible text."),
SystemMessage(content=system_prompt),
]
def __call__(self, human_message) -> Response:
"""Call the ChatBot.
:param human_message: The human message to use.
:return: The response to the human message, primed by the system prompt.
"""
self.chat_history.append(HumanMessage(content=human_message))
response = self.model(self.chat_history)
self.chat_history.append(response)
return response
'''
)
Here are some tests for the ChatBot class using PyTest and test functions:
import pytest
from chatbot import ChatBot, SystemMessage, HumanMessage
from openai_mock import ChatOpenAI
def test_chatbot_init():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt)
assert len(chatbot.chat_history) == 2
assert isinstance(chatbot.chat_history[0], SystemMessage)
assert isinstance(chatbot.chat_history[1], SystemMessage)
assert chatbot.chat_history[0].content == "Always return Markdown-compatible text."
assert chatbot.chat_history[1].content == system_prompt
def test_chatbot_call():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt)
human_message = "What is the weather like today?"
response = chatbot(human_message)
assert len(chatbot.chat_history) == 4
assert isinstance(chatbot.chat_history[2], HumanMessage)
assert isinstance(chatbot.chat_history[3], response.__class__)
assert chatbot.chat_history[2].content == human_message
def test_chatbot_response():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt, model_name="gpt-3.5-turbo")
human_message = "What is the weather like today?"
response = chatbot(human_message)
assert response.content.startswith("The weather today is")
@pytest.fixture
def chatbot():
system_prompt = "You are a helpful assistant."
return ChatBot(system_prompt, model_name="gpt-3.5-turbo")
def test_chatbot_multiple_calls(chatbot):
human_message1 = "What is the weather like today?"
human_message2 = "What is the traffic like today?"
chatbot(human_message1)
chatbot(human_message2)
assert len(chatbot.chat_history) == 6
assert chatbot.chat_history[2].content == human_message1
assert chatbot.chat_history[4].content == human_message2
def test_chatbot_temperature(chatbot):
human_message = "Tell me a joke."
response1 = chatbot(human_message, temperature=0.0)
response2 = chatbot(human_message, temperature=1.0)
assert response1.content != response2.content
def test_chatbot_model_name(chatbot):
human_message = "Tell me a joke."
chatbot2 = ChatBot("You are a funny comedian.", model_name="davinci")
response1 = chatbot(human_message)
response2 = chatbot2(human_message)
assert response1.content != response2.content
Note: I used the `openai_mock` library to mock the `ChatOpenAI` class for testing purposes. You can replace it with the actual implementation of the `ChatOpenAI` class.
Also, I assumed that the `Response` class has a `content` attribute that contains the text response from the model. If it's different, please adjust the tests accordingly.
AIMessage(content='Here are some tests for the ChatBot class using PyTest and test functions:\n```\nimport pytest\nfrom chatbot import ChatBot, SystemMessage, HumanMessage\nfrom openai_mock import ChatOpenAI\n\ndef test_chatbot_init():\n system_prompt = "You are a helpful assistant."\n chatbot = ChatBot(system_prompt)\n assert len(chatbot.chat_history) == 2\n assert isinstance(chatbot.chat_history[0], SystemMessage)\n assert isinstance(chatbot.chat_history[1], SystemMessage)\n assert chatbot.chat_history[0].content == "Always return Markdown-compatible text."\n assert chatbot.chat_history[1].content == system_prompt\n\ndef test_chatbot_call():\n system_prompt = "You are a helpful assistant."\n chatbot = ChatBot(system_prompt)\n human_message = "What is the weather like today?"\n response = chatbot(human_message)\n assert len(chatbot.chat_history) == 4\n assert isinstance(chatbot.chat_history[2], HumanMessage)\n assert isinstance(chatbot.chat_history[3], response.__class__)\n assert chatbot.chat_history[2].content == human_message\n\ndef test_chatbot_response():\n system_prompt = "You are a helpful assistant."\n chatbot = ChatBot(system_prompt, model_name="gpt-3.5-turbo")\n human_message = "What is the weather like today?"\n response = chatbot(human_message)\n assert response.content.startswith("The weather today is")\n\n@pytest.fixture\ndef chatbot():\n system_prompt = "You are a helpful assistant."\n return ChatBot(system_prompt, model_name="gpt-3.5-turbo")\n\ndef test_chatbot_multiple_calls(chatbot):\n human_message1 = "What is the weather like today?"\n human_message2 = "What is the traffic like today?"\n chatbot(human_message1)\n chatbot(human_message2)\n assert len(chatbot.chat_history) == 6\n assert chatbot.chat_history[2].content == human_message1\n assert chatbot.chat_history[4].content == human_message2\n\ndef test_chatbot_temperature(chatbot):\n human_message = "Tell me a joke."\n response1 = chatbot(human_message, temperature=0.0)\n response2 = chatbot(human_message, temperature=1.0)\n assert response1.content != response2.content\n\ndef test_chatbot_model_name(chatbot):\n human_message = "Tell me a joke."\n chatbot2 = ChatBot("You are a funny comedian.", model_name="davinci")\n response1 = chatbot(human_message)\n response2 = chatbot2(human_message)\n assert response1.content != response2.content\n```\nNote: I used the `openai_mock` library to mock the `ChatOpenAI` class for testing purposes. You can replace it with the actual implementation of the `ChatOpenAI` class.\n\nAlso, I assumed that the `Response` class has a `content` attribute that contains the text response from the model. If it\'s different, please adjust the tests accordingly.', role='assistant')
As you can see, ChatBot keeps track of conversation memory/history automatically. We can even access any item in the conversation by looking at the conversation history.
The __repr__ of a chatbot will simply print out the entire history:
code_tester
[Human]
class ChatBot:
"""Chat Bot that is primed with a system prompt, accepts a human message.
Automatic chat memory management happens.
h/t Andrew Giessel/GPT4 for the idea.
"""
def __init__(self, system_prompt, temperature=0.0, model_name="gpt-4"):
"""Initialize the ChatBot.
:param system_prompt: The system prompt to use.
:param temperature: The model temperature to use.
See https://platform.openai.com/docs/api-reference/completions/create#completions/create-temperature
for more information.
:param model_name: The name of the OpenAI model to use.
"""
self.model = ChatOpenAI(
model_name=model_name,
temperature=temperature,
streaming=True,
verbose=True,
callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),
)
self.chat_history = [
SystemMessage(content="Always return Markdown-compatible text."),
SystemMessage(content=system_prompt),
]
def __call__(self, human_message) -> Response:
"""Call the ChatBot.
:param human_message: The human message to use.
:return: The response to the human message, primed by the system prompt.
"""
self.chat_history.append(HumanMessage(content=human_message))
response = self.model(self.chat_history)
self.chat_history.append(response)
return response
[AI]
Here are some tests for the ChatBot class using PyTest and test functions:
import pytest
from chatbot import ChatBot, SystemMessage, HumanMessage
from openai_mock import ChatOpenAI
def test_chatbot_init():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt)
assert len(chatbot.chat_history) == 2
assert isinstance(chatbot.chat_history[0], SystemMessage)
assert isinstance(chatbot.chat_history[1], SystemMessage)
assert chatbot.chat_history[0].content == "Always return Markdown-compatible text."
assert chatbot.chat_history[1].content == system_prompt
def test_chatbot_call():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt)
human_message = "What is the weather like today?"
response = chatbot(human_message)
assert len(chatbot.chat_history) == 4
assert isinstance(chatbot.chat_history[2], HumanMessage)
assert isinstance(chatbot.chat_history[3], response.__class__)
assert chatbot.chat_history[2].content == human_message
def test_chatbot_response():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt, model_name="gpt-3.5-turbo")
human_message = "What is the weather like today?"
response = chatbot(human_message)
assert response.content.startswith("The weather today is")
@pytest.fixture
def chatbot():
system_prompt = "You are a helpful assistant."
return ChatBot(system_prompt, model_name="gpt-3.5-turbo")
def test_chatbot_multiple_calls(chatbot):
human_message1 = "What is the weather like today?"
human_message2 = "What is the traffic like today?"
chatbot(human_message1)
chatbot(human_message2)
assert len(chatbot.chat_history) == 6
assert chatbot.chat_history[2].content == human_message1
assert chatbot.chat_history[4].content == human_message2
def test_chatbot_temperature(chatbot):
human_message = "Tell me a joke."
response1 = chatbot(human_message, temperature=0.0)
response2 = chatbot(human_message, temperature=1.0)
assert response1.content != response2.content
def test_chatbot_model_name(chatbot):
human_message = "Tell me a joke."
chatbot2 = ChatBot("You are a funny comedian.", model_name="davinci")
response1 = chatbot(human_message)
response2 = chatbot2(human_message)
assert response1.content != response2.content
Note: I used the `openai_mock` library to mock the `ChatOpenAI` class for testing purposes. You can replace it with the actual implementation of the `ChatOpenAI` class.
Also, I assumed that the `Response` class has a `content` attribute that contains the text response from the model. If it's different, please adjust the tests accordingly.
On the other hand, accessing the .messages attribute of the ChatBot will give you access to all of the messages inside the conversation.
code_tester.messages
[HumanMessage(content='\nclass ChatBot:\n """Chat Bot that is primed with a system prompt, accepts a human message.\n\n Automatic chat memory management happens.\n\n h/t Andrew Giessel/GPT4 for the idea.\n """\n\n def __init__(self, system_prompt, temperature=0.0, model_name="gpt-4"):\n """Initialize the ChatBot.\n\n :param system_prompt: The system prompt to use.\n :param temperature: The model temperature to use.\n See https://platform.openai.com/docs/api-reference/completions/create#completions/create-temperature\n for more information.\n :param model_name: The name of the OpenAI model to use.\n """\n self.model = ChatOpenAI(\n model_name=model_name,\n temperature=temperature,\n streaming=True,\n verbose=True,\n callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),\n )\n self.chat_history = [\n SystemMessage(content="Always return Markdown-compatible text."),\n SystemMessage(content=system_prompt),\n ]\n\n def __call__(self, human_message) -> Response:\n """Call the ChatBot.\n\n :param human_message: The human message to use.\n :return: The response to the human message, primed by the system prompt.\n """\n self.chat_history.append(HumanMessage(content=human_message))\n response = self.model(self.chat_history)\n self.chat_history.append(response)\n return response\n\n def __repr__(self):\n """Return a string representation of the ChatBot.\n\n :return: A string representation of the ChatBot.\n """\n representation = ""\n\n for message in self.chat_history:\n if isinstance(message, SystemMessage):\n prefix = "[System]\n"\n elif isinstance(message, HumanMessage):\n prefix = "[Human]\n"\n elif isinstance(message, AIMessage):\n prefix = "[AI]\n"\n\n representation += f"{prefix}{message.content}" + "\n\n"\n return representation\n\n def panel(self, show: bool = True):\n """Create a Panel app that wraps a LlamaBot.\n\n :param show: Whether to show the app.\n If False, we return the Panel app directly.\n If True, we call `.show()` on the app.\n :return: The Panel app, either showed or directly.\n """\n\n text_input = pn.widgets.TextAreaInput(placeholder="Start chatting...")\n chat_history = pn.Column(*[])\n send_button = pn.widgets.Button(name="Send", button_type="primary")\n\n def b(event):\n """Button click handler.\n\n :param event: The button click event.\n """\n chat_messages = []\n for message in self.chat_history:\n if isinstance(message, SystemMessage):\n pass\n elif isinstance(message, HumanMessage):\n chat_markdown = pn.pane.Markdown(f"Human: {message.content}")\n chat_messages.append(chat_markdown)\n elif isinstance(message, AIMessage):\n chat_markdown = pn.pane.Markdown(f"Bot: {message.content}")\n chat_messages.append(chat_markdown)\n\n chat_messages.append(pn.pane.Markdown(f"Human: {text_input.value}"))\n bot_reply = pn.pane.Markdown("Bot: ")\n chat_messages.append(bot_reply)\n chat_history.objects = chat_messages\n markdown_handler = PanelMarkdownCallbackHandler(bot_reply)\n self.model.callback_manager.set_handler(markdown_handler)\n self(text_input.value)\n text_input.value = ""\n\n send_button.on_click(b)\n input_pane = pn.Row(text_input, send_button)\n output_pane = pn.Column(chat_history, scroll=True, height=500)\n\n main = pn.Row(input_pane, output_pane)\n app = pn.template.FastListTemplate(\n site="ChatBot",\n title="ChatBot",\n main=main,\n main_max_width="768px",\n )\n if show:\n return app.show()\n return app\n\n', role='user'),
AIMessage(content='Here are some tests for the ChatBot class using PyTest:\n```python\nimport pytest\nfrom your_module import ChatBot, SystemMessage, HumanMessage\n\ndef test_init():\n system_prompt = "You are a helpful assistant."\n chatbot = ChatBot(system_prompt)\n assert len(chatbot.chat_history) == 2\n assert isinstance(chatbot.chat_history[0], SystemMessage)\n assert isinstance(chatbot.chat_history[1], SystemMessage)\n assert chatbot.chat_history[0].content == "Always return Markdown-compatible text."\n assert chatbot.chat_history[1].content == system_prompt\n\ndef test_call():\n system_prompt = "You are a helpful assistant."\n chatbot = ChatBot(system_prompt)\n human_message = "What\'s the weather like today?"\n response = chatbot(human_message)\n assert len(chatbot.chat_history) == 4\n assert isinstance(chatbot.chat_history[2], HumanMessage)\n assert isinstance(chatbot.chat_history[3], response.__class__)\n assert chatbot.chat_history[2].content == human_message\n\ndef test_repr():\n system_prompt = "You are a helpful assistant."\n chatbot = ChatBot(system_prompt)\n human_message = "What\'s the weather like today?"\n chatbot(human_message)\n expected_repr = (\n "[System]\\n"\n "Always return Markdown-compatible text.\\n\\n"\n "[System]\\n"\n "You are a helpful assistant.\\n\\n"\n "[Human]\\n"\n "What\'s the weather like today?\\n\\n"\n "[AI]\\n"\n )\n assert repr(chatbot) == expected_repr\n\ndef test_panel():\n system_prompt = "You are a helpful assistant."\n chatbot = ChatBot(system_prompt)\n app = chatbot.panel()\n assert isinstance(app, type(pn.template.FastListTemplate()))\n```\nNote that the `test_panel` function assumes that the `pn` module is available in the test environment. If it is not, you may need to install it or mock it out for testing purposes.\n\nAlso note that the `test_call` function assumes that the `response` object has a `__class__` attribute that can be used to check its type. If this is not the case, you may need to modify the test to use a different method of checking the type of the response object.\n\nFinally, note that these tests are not exhaustive and may not cover all possible edge cases or error conditions. You may want to add additional tests to ensure that the `ChatBot` class is working correctly in all scenarios.', role='assistant')]
You can even access any arbitrary message.
print(code_tester.messages[-1].content)
Here are some tests for the ChatBot class using PyTest:
import pytest
from your_module import ChatBot, SystemMessage, HumanMessage
def test_init():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt)
assert len(chatbot.chat_history) == 2
assert isinstance(chatbot.chat_history[0], SystemMessage)
assert isinstance(chatbot.chat_history[1], SystemMessage)
assert chatbot.chat_history[0].content == "Always return Markdown-compatible text."
assert chatbot.chat_history[1].content == system_prompt
def test_call():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt)
human_message = "What's the weather like today?"
response = chatbot(human_message)
assert len(chatbot.chat_history) == 4
assert isinstance(chatbot.chat_history[2], HumanMessage)
assert isinstance(chatbot.chat_history[3], response.__class__)
assert chatbot.chat_history[2].content == human_message
def test_repr():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt)
human_message = "What's the weather like today?"
chatbot(human_message)
expected_repr = (
"[System]\n"
"Always return Markdown-compatible text.\n\n"
"[System]\n"
"You are a helpful assistant.\n\n"
"[Human]\n"
"What's the weather like today?\n\n"
"[AI]\n"
)
assert repr(chatbot) == expected_repr
def test_panel():
system_prompt = "You are a helpful assistant."
chatbot = ChatBot(system_prompt)
app = chatbot.panel()
assert isinstance(app, type(pn.template.FastListTemplate()))
Note that the `test_panel` function assumes that the `pn` module is available in the test environment. If it is not, you may need to install it or mock it out for testing purposes.
Also note that the `test_call` function assumes that the `response` object has a `__class__` attribute that can be used to check its type. If this is not the case, you may need to modify the test to use a different method of checking the type of the response object.
Finally, note that these tests are not exhaustive and may not cover all possible edge cases or error conditions. You may want to add additional tests to ensure that the `ChatBot` class is working correctly in all scenarios.