mirror of
https://github.com/hwchase17/langchain
synced 2024-11-08 07:10:35 +00:00
a7af32c274
# Cassandra support for chat history ### Description - Store chat messages in cassandra ### Dependency - cassandra-driver - Python Module ## Before submitting - Added Integration Test ## Who can review? @hwchase17 @agola11 # Your PR Title (What it does) <!-- Thank you for contributing to LangChain! Your PR will appear in our next release under the title you set. Please make sure it highlights your valuable contribution. Replace this with a description of the change, the issue it fixes (if applicable), and relevant context. List any dependencies required for this change. After you're done, someone will review your PR. They may suggest improvements. If no one reviews your PR within a few days, feel free to @-mention the same people again, as notifications can get lost. --> <!-- Remove if not applicable --> Fixes # (issue) ## Before submitting <!-- If you're adding a new integration, include an integration test and an example notebook showing its use! --> ## Who can review? Community members can review the PR once tests pass. Tag maintainers/contributors who might be interested: <!-- For a quicker response, figure out the right person to tag with @ @hwchase17 - project lead Tracing / Callbacks - @agola11 Async - @agola11 DataLoaders - @eyurtsev Models - @hwchase17 - @agola11 Agents / Tools / Toolkits - @vowelparrot VectorStores / Retrievers / Memory - @dev2049 --> Co-authored-by: Jinto Jose <129657162+jj701@users.noreply.github.com>
43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
import json
|
|
import os
|
|
|
|
from langchain.memory import ConversationBufferMemory
|
|
from langchain.memory.chat_message_histories.cassandra import (
|
|
CassandraChatMessageHistory,
|
|
)
|
|
from langchain.schema import _message_to_dict
|
|
|
|
# Replace these with your cassandra contact points
|
|
contact_points = (
|
|
os.environ["CONTACT_POINTS"].split(",")
|
|
if "CONTACT_POINTS" in os.environ
|
|
else ["cassandra"]
|
|
)
|
|
|
|
|
|
def test_memory_with_message_store() -> None:
|
|
"""Test the memory with a message store."""
|
|
# setup cassandra as a message store
|
|
message_history = CassandraChatMessageHistory(
|
|
contact_points=contact_points, session_id="test-session"
|
|
)
|
|
memory = ConversationBufferMemory(
|
|
memory_key="baz", chat_memory=message_history, return_messages=True
|
|
)
|
|
|
|
# add some messages
|
|
memory.chat_memory.add_ai_message("This is me, the AI")
|
|
memory.chat_memory.add_user_message("This is me, the human")
|
|
|
|
# get the message history from the memory store and turn it into a json
|
|
messages = memory.chat_memory.messages
|
|
messages_json = json.dumps([_message_to_dict(msg) for msg in messages])
|
|
|
|
assert "This is me, the AI" in messages_json
|
|
assert "This is me, the human" in messages_json
|
|
|
|
# remove the record from Cassandra, so the next test run won't pick it up
|
|
memory.chat_memory.clear()
|
|
|
|
assert memory.chat_memory.messages == []
|