AWS DynamoDB

Amazon AWS DynamoDB is a fully managed NoSQL database service that provides fast and predictable performance with seamless scalability.

This notebook goes over how to use DynamoDB to store chat message history.

First make sure you have correctly configured the AWS CLI. Then make sure you have installed boto3.

!pip install boto3

Next, create the DynamoDB Table where we will be storing messages:

import boto3

# Get the service resource.
dynamodb = boto3.resource("dynamodb")

# Create the DynamoDB table.
table = dynamodb.create_table(
    TableName="SessionTable",
    KeySchema=[{"AttributeName": "SessionId", "KeyType": "HASH"}],
    AttributeDefinitions=[{"AttributeName": "SessionId", "AttributeType": "S"}],
    BillingMode="PAY_PER_REQUEST",
)

# Wait until the table exists.
table.meta.client.get_waiter("table_exists").wait(TableName="SessionTable")

# Print out some data about the table.
print(table.item_count)

0

DynamoDBChatMessageHistory

from langchain.memory.chat_message_histories import DynamoDBChatMessageHistory

history = DynamoDBChatMessageHistory(table_name="SessionTable", session_id="0")

history.add_user_message("hi!")

history.add_ai_message("whats up?")

history.messages

[HumanMessage(content='hi!', additional_kwargs={}, example=False),
 AIMessage(content='whats up?', additional_kwargs={}, example=False),
 HumanMessage(content='hi!', additional_kwargs={}, example=False),
 AIMessage(content='whats up?', additional_kwargs={}, example=False)]

DynamoDBChatMessageHistory with Custom Endpoint URL

Sometimes it is useful to specify the URL to the AWS endpoint to connect to. For instance, when you are running locally against Localstack. For those cases you can specify the URL via the endpoint_url parameter in the constructor.

from langchain.memory.chat_message_histories import DynamoDBChatMessageHistory

history = DynamoDBChatMessageHistory(
    table_name="SessionTable",
    session_id="0",
    endpoint_url="http://localhost.localstack.cloud:4566",
)

DynamoDBChatMessageHistory With Different Keys Composite Keys

The default key for DynamoDBChatMessageHistory is {"SessionId": self.session_id}, but you can modify this to match your table design.

Primary Key Name

You may modify the primary key by passing in a primary_key_name value in the constructor, resulting in the following: {self.primary_key_name: self.session_id}

Composite Keys

When using an existing DynamoDB table, you may need to modify the key structure from the default of to something including a Sort Key. To do this you may use the key parameter.

Passing a value for key will override the primary_key parameter, and the resulting key structure will be the passed value.

from langchain.memory.chat_message_histories import DynamoDBChatMessageHistory

composite_table = dynamodb.create_table(
    TableName="CompositeTable",
    KeySchema=[
        {"AttributeName": "PK", "KeyType": "HASH"},
        {"AttributeName": "SK", "KeyType": "RANGE"},
    ],
    AttributeDefinitions=[
        {"AttributeName": "PK", "AttributeType": "S"},
        {"AttributeName": "SK", "AttributeType": "S"},
    ],
    BillingMode="PAY_PER_REQUEST",
)

# Wait until the table exists.
composite_table.meta.client.get_waiter("table_exists").wait(TableName="CompositeTable")

# Print out some data about the table.
print(composite_table.item_count)

my_key = {
    "PK": "session_id::0",
    "SK": "langchain_history",
}

composite_key_history = DynamoDBChatMessageHistory(
    table_name="CompositeTable",
    session_id="0",
    endpoint_url="http://localhost.localstack.cloud:4566",
    key=my_key,
)

composite_key_history.add_user_message("hello, composite dynamodb table!")

composite_key_history.messages

0

[HumanMessage(content='hello, composite dynamodb table!', additional_kwargs={}, example=False)]

Agent with DynamoDB Memory

from langchain.agents import AgentType, Tool, initialize_agent
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain_experimental.utilities import PythonREPL

message_history = DynamoDBChatMessageHistory(table_name="SessionTable", session_id="1")
memory = ConversationBufferMemory(
    memory_key="chat_history", chat_memory=message_history, return_messages=True
)

python_repl = PythonREPL()

# You can create the tool to pass to an agent
tools = [
    Tool(
        name="python_repl",
        description="A Python shell. Use this to execute python commands. Input should be a valid python command. If you want to see the output of a value, you should print it out with `print(...)`.",
        func=python_repl.run,
    )
]

llm = ChatOpenAI(temperature=0)
agent_chain = initialize_agent(
    tools,
    llm,
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
    verbose=True,
    memory=memory,
)

agent_chain.run(input="Hello!")

agent_chain.run(input="Who owns Twitter?")

agent_chain.run(input="My name is Bob.")

agent_chain.run(input="Who am I?")