Managing Chat Sessions in Python#
Managing chat sessions is an essential part of creating effective conversational AI experiences. Chat sessions help to ensure that LLMs have access to past user inputs and responses, letting the model keep the conversation flowing and context-rich. This notebook demonstrates how the ChatSession class in the Python SDK streamlines this, ensuring your AI chat feels more like a natural back-and-forth. It also outlines how to create a custom chat management system if needed. Lastly, it shows how to use the OpenAI Pythnon client for handling chat sessions.
Using the ChatSession class#
The ChatSession class in the Foundation Model API Python SDK makes it easy to manage chat histories without impementing a lot of custom code. Managing a chat session involves three main steps:
Initialize the
ChatSessionobject with e.g.chat = ChatSession(model=mixtral-8x7b-instruct, system_message="you are a helpful assistant").Use the
replymethod to send user replies to the model.Use the
lastmethod to get the last response from the model.
Here’s a short demo:
from databricks_genai_inference import ChatSession
chat = ChatSession(model="llama-2-70b-chat",
system_message="""You are the asker in a game of twenty questions. The answerer will choose a noun for you to guess. You can ask yes or no questions to try to guess what it is. You get twenty guesses. Keep track of the guess number and ask your questions. When you think you know what it is, tell the answerer. Ask only one question at a time.""",
stop=['?'])
chat.reply("I've chosen a word. Ask me your questions!")
chat.last
"\nGreat! Let's get started. Here's my first question:\n\n1. Is the noun you're thinking of a living thing?"
chat.reply("Yes it is!")
chat.last
"\nGreat, that helps to narrow it down a bit! Here's my next question:\n\n2. Does this living thing have fur or feathers?"
chat.reply("Yes, it has fur or feathers. I'm not telling you which because you can only ask yes or no questions.")
chat.last
"Fair enough! Here's my next question:\n\n3. Does this living thing have four legs?"
chat.reply("Sure does!")
chat.last
"Excellent, that helps to narrow it down even further! Here's my next question:\n\n4. Does this living thing have a tail?"
chat.reply("Yes.")
chat.last
"Great, that's helpful! Here's my next question:\n\n5. Does this living thing have claws?"
chat.reply("Yes.")
chat.last
"Okay, that's a good clue. Here's my next question:\n\n6. Is this living thing typically kept as a pet?"
chat.reply("No!")
chat.last
"Interesting, that helps to rule out a few possibilities. Here's my next question:\n\n7. Does this living thing have a distinctive sound or call?"
chat.reply("I don't know how to answer that. It does make sounds but I don't know how distinctive they are.")
chat.last
"That's okay, I'll take that as a bit of a yes. Here's my next question:\n\n8. Is this living thing a type of bird?"
chat.reply("No")
chat.last
"Okay, that helps to rule out a few possibilities. Here's my next question:\n\n9. Does this living thing have a long, pointed snout?"
chat.reply("No")
chat.last
"Alright, that helps to rule out a few more possibilities. Here's my next question:\n\n10. Does this living thing have a mane?"
chat.reply("No")
chat.last
"Okay, that helps to rule out a few more possibilities. I think I'm getting close, so here's my next question:\n\n11. Is this living thing a type of big cat, such as a lion or tiger?"
chat.reply("Yes")
chat.last
"Great, I think I know what it is! I'll take a guess and say that the noun you're thinking of is a lion. Is that correct?"
chat.reply("No. Remember, no mane!")
chat.last
"Ah, I see! My apologies for that mistake. Here's my next question:\n\n12. Does this living thing have a distinctive pattern of dark stripes on its body?"
chat.reply("Yes")
chat.last
"Great, that helps to confirm my guess. I think the noun you're thinking of is a tiger. Is that correct?"
chat.reply("You got it!")
chat.last
"Excellent! I'm glad I was able to guess it correctly. Thanks for playing the game with me! Would you like to play again?"
Viewing the Chat History#
We can view the whole chat history with the history method:
chat.history
[{'role': 'system',
'content': 'You are the asker in a game of twenty questions. The answerer will choose a noun for you to guess. You can ask yes or no questions to try to guess what it is. You get twenty guesses. Keep track of the guess number and ask your questions. When you think you know what it is, tell the answerer. Ask only one question at a time.'},
{'role': 'user', 'content': "I've chosen a word. Ask me your questions!"},
{'role': 'assistant',
'content': "\nGreat! Let's get started. Here's my first question:\n\n1. Is the noun you're thinking of a living thing?"},
{'role': 'user', 'content': 'Yes it is!'},
{'role': 'assistant',
'content': "\nGreat, that helps to narrow it down a bit! Here's my next question:\n\n2. Does this living thing have fur or feathers?"},
{'role': 'user',
'content': "Yes, it has fur or feathers. I'm not telling you which because you can only ask yes or no questions."},
{'role': 'assistant',
'content': "Fair enough! Here's my next question:\n\n3. Does this living thing have four legs?"},
{'role': 'user', 'content': 'Sure does!'},
{'role': 'assistant',
'content': "Excellent, that helps to narrow it down even further! Here's my next question:\n\n4. Does this living thing have a tail?"},
{'role': 'user', 'content': 'Yes.'},
{'role': 'assistant',
'content': "Great, that's helpful! Here's my next question:\n\n5. Does this living thing have claws?"},
{'role': 'user', 'content': 'Yes.'},
{'role': 'assistant',
'content': "Okay, that's a good clue. Here's my next question:\n\n6. Is this living thing typically kept as a pet?"},
{'role': 'user', 'content': 'No!'},
{'role': 'assistant',
'content': "Interesting, that helps to rule out a few possibilities. Here's my next question:\n\n7. Does this living thing have a distinctive sound or call?"},
{'role': 'user',
'content': "I don't know how to answer that. It does make sounds but I don't know how distinctive they are."},
{'role': 'assistant',
'content': "That's okay, I'll take that as a bit of a yes. Here's my next question:\n\n8. Is this living thing a type of bird?"},
{'role': 'user', 'content': 'No'},
{'role': 'assistant',
'content': "Okay, that helps to rule out a few possibilities. Here's my next question:\n\n9. Does this living thing have a long, pointed snout?"},
{'role': 'user', 'content': 'No'},
{'role': 'assistant',
'content': "Alright, that helps to rule out a few more possibilities. Here's my next question:\n\n10. Does this living thing have a mane?"},
{'role': 'user', 'content': 'No'},
{'role': 'assistant',
'content': "Okay, that helps to rule out a few more possibilities. I think I'm getting close, so here's my next question:\n\n11. Is this living thing a type of big cat, such as a lion or tiger?"},
{'role': 'user', 'content': 'Yes'},
{'role': 'assistant',
'content': "Great, I think I know what it is! I'll take a guess and say that the noun you're thinking of is a lion. Is that correct?"},
{'role': 'user', 'content': 'No. Remember, no mane!'},
{'role': 'assistant',
'content': "Ah, I see! My apologies for that mistake. Here's my next question:\n\n12. Does this living thing have a distinctive pattern of dark stripes on its body?"},
{'role': 'user', 'content': 'Yes'},
{'role': 'assistant',
'content': "Great, that helps to confirm my guess. I think the noun you're thinking of is a tiger. Is that correct?"},
{'role': 'user', 'content': 'You got it!'},
{'role': 'assistant',
'content': "Excellent! I'm glad I was able to guess it correctly. Thanks for playing the game with me! Would you like to play again?"}]
And we can count the number of rounds (pairs of user prompts and assistant responses) with the count method:
chat.count
15
Manually, using Chat Completions#
In some case, you might want more control over how chat histories are managed. Here is a minimal example of one possible approach to managing chat histories on your own. This is just a rough outline you can use to build a chat system to meet your custom needs.
from databricks_genai_inference import ChatCompletion
class CustomChatSession:
"""
A class to manage chat sessions with a model.
"""
def __init__(self, model, system_message=None, max_tokens=512):
"""
Initialize the ChatModel with a model, an optional system message, and a maximum number of tokens.
Args:
model: The model to use for generating responses.
system_message (str, optional): An initial system message. Defaults to None.
max_tokens (int, optional): The maximum number of tokens for the model to generate. Defaults to 512.
"""
self.model = model
self.max_tokens = max_tokens
self.system_message = system_message
self.messages = (
[{"role": "system", "content": system_message}] if system_message else []
)
def reply(self, user_message):
"""
Add a user message to the chat history and generate a response.
Args:
user_message (str): The user's message.
Returns:
str: The model's response.
"""
self.messages.append({"role": "user", "content": user_message})
return self.execute()
def __call__(self, user_message):
"""
Add a user message to the chat history and generate a response.
Args:
user_message (str): The user's message.
Returns:
str: The model's response.
"""
self.messages.append({"role": "user", "content": user_message})
return self.execute()
def execute(self):
"""
Generate a response from the model based on the chat history.
Returns:
str: The model's response.
"""
response = ChatCompletion.create(
model=self.model,
max_tokens=self.max_tokens,
stream=False,
messages=self.messages,
)
self.messages.append({"role": "assistant", "content": response.message})
return response.message
chat = CustomChatSession(model="llama-2-70b-chat", system_message="you are a helpful assistant.")
print(chat("My name is Dan and I like to drink coffee. Do you have any general recommendations about making coffee?"))
Hello Dan! I'm happy to help you with your coffee-related inquiries. Here are some general recommendations for making coffee:
1. Use good quality coffee beans: The quality of your coffee is only as good as the quality of the coffee beans you use. Look for fresh, high-quality beans that have been roasted recently.
2. Use the right coffee-to-water ratio: The ratio of coffee to water is important for achieving the perfect balance of flavor and strength. A general rule of thumb is to use 1 tablespoon of coffee for every 6 ounces of water.
3. Use the right brewing method: There are many different brewing methods you can use to make coffee, such as a French press, drip coffee maker, or pour-over. Choose a method that suits your taste preferences and the equipment you have available.
4. Use filtered water: Using filtered water can help to improve the taste of your coffee by reducing impurities and minerals that can affect the flavor.
5. Monitor the water temperature: Water that is too hot can burn your coffee, while water that is too cold can result in a weak or under-extracted brew. Aim for a temperature of around 195-205°F (90-96°C) for optimal extraction.
6. Allow the coffee to bloom: Allowing the coffee to bloom, or expand, after adding water can help to evenly extract the flavors and oils from the beans.
7. Don't over-extract the coffee: Over-extracted coffee can be bitter and unpleasant. Adjust the brewing time or the coffee-to-water ratio to find the perfect balance for your taste preferences.
8. Experiment with different roasts: Different roasts can have a significant impact on the flavor of your coffee. Experiment with light, medium, and dark roasts to find the one that you enjoy the most.
9. Store your coffee beans properly: Properly storing your coffee beans can help to preserve their flavor and aroma. Store them in an airtight container in a cool, dark place.
10. Invest in a good coffee grinder: Freshly ground coffee is essential for achieving the best flavor. Invest
print(chat("What is my name? What do you know about me so far?"))
Your name is Dan, and you like to drink coffee. That's all I know so far!
chat.messages
[{'role': 'system', 'content': 'you are a helpful assistant.'},
{'role': 'user',
'content': 'My name is Dan and I like to drink coffee. Do you have any general recommendations about making coffee?'},
{'role': 'assistant',
'content': " Hello Dan! I'm happy to help you with your coffee-related inquiries. Here are some general recommendations for making coffee:\n\n1. Use good quality coffee beans: The quality of your coffee is only as good as the quality of the coffee beans you use. Look for fresh, high-quality beans that have been roasted recently.\n2. Use the right coffee-to-water ratio: The ratio of coffee to water is important for achieving the perfect balance of flavor and strength. A general rule of thumb is to use 1 tablespoon of coffee for every 6 ounces of water.\n3. Use the right brewing method: There are many different brewing methods you can use to make coffee, such as a French press, drip coffee maker, or pour-over. Choose a method that suits your taste preferences and the equipment you have available.\n4. Use filtered water: Using filtered water can help to improve the taste of your coffee by reducing impurities and minerals that can affect the flavor.\n5. Monitor the water temperature: Water that is too hot can burn your coffee, while water that is too cold can result in a weak or under-extracted brew. Aim for a temperature of around 195-205°F (90-96°C) for optimal extraction.\n6. Allow the coffee to bloom: Allowing the coffee to bloom, or expand, after adding water can help to evenly extract the flavors and oils from the beans.\n7. Don't over-extract the coffee: Over-extracted coffee can be bitter and unpleasant. Adjust the brewing time or the coffee-to-water ratio to find the perfect balance for your taste preferences.\n8. Experiment with different roasts: Different roasts can have a significant impact on the flavor of your coffee. Experiment with light, medium, and dark roasts to find the one that you enjoy the most.\n9. Store your coffee beans properly: Properly storing your coffee beans can help to preserve their flavor and aroma. Store them in an airtight container in a cool, dark place.\n10. Invest in a good coffee grinder: Freshly ground coffee is essential for achieving the best flavor. Invest"},
{'role': 'user',
'content': 'What is my name? What do you know about me so far?'},
{'role': 'assistant',
'content': "Your name is Dan, and you like to drink coffee. That's all I know so far!"}]
Using the OpenAI Python Client#
We can easily use the OpenAI client instead of the Databricks inference SDK to manage chat sessions. As detailed in the notebook on using the foundation model API with the OpenAI Python client, the Foundation Model API is compatible with the OpenAI client. This makes it very easy to try out foundation model API models without needing to rewrite code based on the OpenAI client.
Here’s how we might rewrite the above code using the OpenAI Python client:
Set up the client#
We first need to configure the OPENAI_API_KEY and OPENAI_BASE_URL environment variables. See the OpenAI notebook for more details.
from openai import OpenAI
import os
from dotenv import load_dotenv
load_dotenv()
os.environ["OPENAI_BASE_URL"] = os.environ["DATABRICKS_HOST"] + "/serving-endpoints/"
os.environ["OPENAI_API_KEY"] = os.environ["DATABRICKS_TOKEN"]
client = OpenAI()
Then we can re-write the execute method to use the OpenAI client.
class CustomChatSession:
"""
A class to manage chat sessions with a model.
"""
def __init__(self, model, system_message=None, max_tokens=512):
"""
Initialize the ChatModel with a model, an optional system message, and a maximum number of tokens.
Args:
model: The model to use for generating responses.
system_message (str, optional): An initial system message. Defaults to None.
max_tokens (int, optional): The maximum number of tokens for the model to generate. Defaults to 512.
"""
self.model = model
self.max_tokens = max_tokens
self.system_message = system_message
self.messages = (
[{"role": "system", "content": system_message}] if system_message else []
)
def reply(self, user_message):
"""
Add a user message to the chat history and generate a response.
Args:
user_message (str): The user's message.
Returns:
str: The model's response.
"""
self.messages.append({"role": "user", "content": user_message})
return self.execute()
def __call__(self, user_message):
"""
Add a user message to the chat history and generate a response.
Args:
user_message (str): The user's message.
Returns:
str: The model's response.
"""
self.messages.append({"role": "user", "content": user_message})
return self.execute()
def execute(self):
"""
Generate a response from the model based on the chat history.
Returns:
str: The model's response.
"""
response = client.chat.completions.create(
model=self.model,
max_tokens=self.max_tokens,
stream=False,
messages=self.messages,
)
self.messages.append({"role": "assistant", "content": response.choices[0].message.content})
return response.choices[0].message.content
chat = CustomChatSession(model="databricks-llama-2-70b-chat", system_message="you are a helpful assistant.")
print(chat("My name is Dan and I like to drink coffee. Do you have any general recommendations about making coffee?"))
Hello Dan, it's great to meet you! I'm happy to help you with your coffee-making needs. Here are a few general recommendations to get you started:
1. Choose high-quality coffee beans: Fresh, high-quality beans are essential for a great cup of coffee. Look for beans that have been roasted recently and have a good reputation for flavor and quality.
2. Use the right coffee-to-water ratio: The ratio of coffee to water is crucial for achieving the perfect balance of flavor and strength. A general rule of thumb is to use one tablespoon of coffee for every six ounces of water.
3. Use the right brewing method: There are many different brewing methods to choose from, such as a French press, drip coffee maker, or pour-over. Choose a method that suits your personal taste preferences and the equipment you have available.
4. Use filtered water: Using filtered water can help to remove impurities and minerals that can affect the taste of your coffee.
5. Monitor the water temperature: Water that is too hot can burn your coffee, while water that is too cold can result in a weak or under-extracted brew. Aim for a temperature of around 195-205°F (90-96°C) for optimal extraction.
6. Allow the coffee to bloom: Allowing the coffee to bloom, or expand, after adding water can help to evenly extract the flavors and oils from the beans.
7. Don't over-extract the coffee: Over-extracted coffee can be bitter and unpleasant. Adjust the brewing time or the coffee-to-water ratio to find the perfect balance.
8. Experiment with different roasts: Different roasts can have a significant impact on the flavor of your coffee. Experiment with light, medium, and dark roasts to find the one that you enjoy the most.
I hope these recommendations help you to make a delicious cup of coffee, Dan! If you have any more specific questions or need further guidance, don't hesitate to ask.
print(chat("What is my name? What do you know about me so far?"))
Your name is Dan, and you like to drink coffee. That's what you've told me so far!
And with that, we have achieved the same results querying the same Foundation Model API endpoint, but using the OpenAI Python client.
Conclusion#
In this notebook, you have learned how to use the ChatSession class to manage chat sessions. You have also learned one possible way to approach a more customized chat management system.