Skip to content

Commit

Permalink
CrateDB memory: Add conversational memory support
Browse files Browse the repository at this point in the history 
The implementation is based on the generic `SQLChatMessageHistory`.
  • Loading branch information
@amotl
amotl committed Sep 25, 2023
1 parent 571ce1f commit 1843ce4
Show file tree
Hide file tree
Showing 6 changed files with 679 additions and 2 deletions.
8 changes: 8 additions & 0 deletions docs/docs_skeleton/vercel.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2632,6 +2632,14 @@
"source": "/docs/modules/memory/integrations/cassandra_chat_message_history",
"destination": "/docs/integrations/memory/cassandra_chat_message_history"
},
{
"source": "/en/latest/modules/memory/examples/cratedb_chat_message_history.html",
"destination": "/docs/integrations/memory/cratedb_chat_message_history"
},
{
"source": "/docs/modules/memory/integrations/cratedb_chat_message_history",
"destination": "/docs/integrations/memory/cratedb_chat_message_history"
},
{
"source": "/en/latest/modules/memory/examples/dynamodb_chat_message_history.html",
"destination": "/docs/integrations/memory/dynamodb_chat_message_history"
Expand Down
357 changes: 357 additions & 0 deletions docs/extras/integrations/memory/cratedb_chat_message_history.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,357 @@
{
"cells": [
{
"cell_type": "markdown",
"source": [
"# CrateDB Chat Message History\n",
"\n",
"This notebook demonstrates how to use the `CrateDBChatMessageHistory`\n",
"to manage chat history in CrateDB, for supporting conversational memory."
],
"metadata": {
"collapsed": false
},
"id": "f22eab3f84cbeb37"
},
{
"cell_type": "markdown",
"source": [
"## Prerequisites"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"!#pip install 'crate[sqlalchemy]' 'langchain'"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "markdown",
"source": [
"## Configuration\n",
"\n",
"To use the storage wrapper, you will need to configure two details.\n",
"\n",
"1. Session Id - a unique identifier of the session, like user name, email, chat id etc.\n",
"2. Database connection string: An SQLAlchemy-compatible URI that specifies the database\n",
" connection. It will be passed to SQLAlchemy create_engine function."
],
"metadata": {
"collapsed": false
},
"id": "f8f2830ee9ca1e01"
},
{
"cell_type": "code",
"execution_count": 52,
"outputs": [],
"source": [
"from langchain.memory.chat_message_histories import CrateDBChatMessageHistory\n",
"\n",
"CONNECTION_STRING = \"crate://crate@localhost:4200/?schema=example\"\n",
"\n",
"chat_message_history = CrateDBChatMessageHistory(\n",
"\tsession_id=\"test_session\",\n",
"\tconnection_string=CONNECTION_STRING\n",
")"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "markdown",
"source": [
"## Basic Usage"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": 53,
"outputs": [],
"source": [
"chat_message_history.add_user_message(\"Hello\")\n",
"chat_message_history.add_ai_message(\"Hi\")"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2023-08-28T10:04:38.077748Z",
"start_time": "2023-08-28T10:04:36.105894Z"
}
},
"id": "4576e914a866fb40"
},
{
"cell_type": "code",
"execution_count": 61,
"outputs": [
{
"data": {
"text/plain": "[HumanMessage(content='Hello', additional_kwargs={}, example=False),\n AIMessage(content='Hi', additional_kwargs={}, example=False)]"
},
"execution_count": 61,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"chat_message_history.messages"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2023-08-28T10:04:38.929396Z",
"start_time": "2023-08-28T10:04:38.915727Z"
}
},
"id": "b476688cbb32ba90"
},
{
"cell_type": "markdown",
"source": [
"## Custom Storage Model\n",
"\n",
"The default data model, which stores information about conversation messages only\n",
"has two slots for storing message details, the session id, and the message dictionary.\n",
"\n",
"If you want to store additional information, like message date, author, language etc.,\n",
"please provide an implementation for a custom message converter.\n",
"\n",
"This example demonstrates how to create a custom message converter, by implementing\n",
"the `BaseMessageConverter` interface."
],
"metadata": {
"collapsed": false
},
"id": "2e5337719d5614fd"
},
{
"cell_type": "code",
"execution_count": 55,
"outputs": [],
"source": [
"from datetime import datetime\n",
"from typing import Any\n",
"\n",
"from langchain.memory.chat_message_histories.cratedb import generate_autoincrement_identifier\n",
"from langchain.memory.chat_message_histories.sql import BaseMessageConverter\n",
"from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage\n",
"\n",
"import sqlalchemy as sa\n",
"from sqlalchemy.orm import declarative_base\n",
"\n",
"\n",
"Base = declarative_base()\n",
"\n",
"\n",
"class CustomMessage(Base):\n",
"\t__tablename__ = \"custom_message_store\"\n",
"\n",
"\tid = sa.Column(sa.BigInteger, primary_key=True, default=generate_autoincrement_identifier)\n",
"\tsession_id = sa.Column(sa.Text)\n",
"\ttype = sa.Column(sa.Text)\n",
"\tcontent = sa.Column(sa.Text)\n",
"\tcreated_at = sa.Column(sa.DateTime)\n",
"\tauthor_email = sa.Column(sa.Text)\n",
"\n",
"\n",
"class CustomMessageConverter(BaseMessageConverter):\n",
"\tdef __init__(self, author_email: str):\n",
"\t\tself.author_email = author_email\n",
"\t\n",
"\tdef from_sql_model(self, sql_message: Any) -> BaseMessage:\n",
"\t\tif sql_message.type == \"human\":\n",
"\t\t\treturn HumanMessage(\n",
"\t\t\t\tcontent=sql_message.content,\n",
"\t\t\t)\n",
"\t\telif sql_message.type == \"ai\":\n",
"\t\t\treturn AIMessage(\n",
"\t\t\t\tcontent=sql_message.content,\n",
"\t\t\t)\n",
"\t\telif sql_message.type == \"system\":\n",
"\t\t\treturn SystemMessage(\n",
"\t\t\t\tcontent=sql_message.content,\n",
"\t\t\t)\n",
"\t\telse:\n",
"\t\t\traise ValueError(f\"Unknown message type: {sql_message.type}\")\n",
"\t\n",
"\tdef to_sql_model(self, message: BaseMessage, session_id: str) -> Any:\n",
"\t\tnow = datetime.now()\n",
"\t\treturn CustomMessage(\n",
"\t\t\tsession_id=session_id,\n",
"\t\t\ttype=message.type,\n",
"\t\t\tcontent=message.content,\n",
"\t\t\tcreated_at=now,\n",
"\t\t\tauthor_email=self.author_email\n",
"\t\t)\n",
"\t\n",
"\tdef get_sql_model_class(self) -> Any:\n",
"\t\treturn CustomMessage\n",
"\n",
"\n",
"if __name__ == \"__main__\":\n",
"\n",
"\tBase.metadata.drop_all(bind=sa.create_engine(CONNECTION_STRING))\n",
"\n",
"\tchat_message_history = CrateDBChatMessageHistory(\n",
"\t\tsession_id=\"test_session\",\n",
"\t\tconnection_string=CONNECTION_STRING,\n",
"\t\tcustom_message_converter=CustomMessageConverter(\n",
"\t\t\tauthor_email=\"test@example.com\"\n",
"\t\t)\n",
"\t)\n",
"\n",
"\tchat_message_history.add_user_message(\"Hello\")\n",
"\tchat_message_history.add_ai_message(\"Hi\")"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2023-08-28T10:04:41.510498Z",
"start_time": "2023-08-28T10:04:41.494912Z"
}
},
"id": "fdfde84c07d071bb"
},
{
"cell_type": "code",
"execution_count": 60,
"outputs": [
{
"data": {
"text/plain": "[HumanMessage(content='Hello', additional_kwargs={}, example=False),\n AIMessage(content='Hi', additional_kwargs={}, example=False)]"
},
"execution_count": 60,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"chat_message_history.messages"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2023-08-28T10:04:43.497990Z",
"start_time": "2023-08-28T10:04:43.492517Z"
}
},
"id": "4a6a54d8a9e2856f"
},
{
"cell_type": "markdown",
"source": [
"## Custom Name for Session Column\n",
"\n",
"The session id, a unique token identifying the session, is an important property of\n",
"this subsystem. If your database table stores it in a different column, you can use\n",
"the `session_id_field_name` keyword argument to adjust the name correspondingly."
],
"metadata": {
"collapsed": false
},
"id": "622aded629a1adeb"
},
{
"cell_type": "code",
"execution_count": 57,
"outputs": [],
"source": [
"import json\n",
"import typing as t\n",
"\n",
"from langchain.memory.chat_message_histories.cratedb import generate_autoincrement_identifier, CrateDBMessageConverter\n",
"from langchain.schema import _message_to_dict\n",
"\n",
"\n",
"Base = declarative_base()\n",
"\n",
"class MessageWithDifferentSessionIdColumn(Base):\n",
"\t__tablename__ = \"message_store_different_session_id\"\n",
"\tid = sa.Column(sa.BigInteger, primary_key=True, default=generate_autoincrement_identifier)\n",
"\tcustom_session_id = sa.Column(sa.Text)\n",
"\tmessage = sa.Column(sa.Text)\n",
"\n",
"\n",
"class CustomMessageConverterWithDifferentSessionIdColumn(CrateDBMessageConverter):\n",
" def __init__(self):\n",
" self.model_class = MessageWithDifferentSessionIdColumn\n",
"\n",
" def to_sql_model(self, message: BaseMessage, custom_session_id: str) -> t.Any:\n",
" return self.model_class(\n",
" custom_session_id=custom_session_id, message=json.dumps(_message_to_dict(message))\n",
" )\n",
"\n",
"\n",
"if __name__ == \"__main__\":\n",
"\tBase.metadata.drop_all(bind=sa.create_engine(CONNECTION_STRING))\n",
"\n",
"\tchat_message_history = CrateDBChatMessageHistory(\n",
"\t\tsession_id=\"test_session\",\n",
"\t\tconnection_string=CONNECTION_STRING,\n",
"\t\tcustom_message_converter=CustomMessageConverterWithDifferentSessionIdColumn(),\n",
"\t\tsession_id_field_name=\"custom_session_id\",\n",
"\t)\n",
"\n",
"\tchat_message_history.add_user_message(\"Hello\")\n",
"\tchat_message_history.add_ai_message(\"Hi\")"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": 58,
"outputs": [
{
"data": {
"text/plain": "[HumanMessage(content='Hello', additional_kwargs={}, example=False),\n AIMessage(content='Hi', additional_kwargs={}, example=False)]"
},
"execution_count": 58,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"chat_message_history.messages"
],
"metadata": {
"collapsed": false
}
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 2
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython2",
"version": "2.7.6"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Loading

0 comments on commit 1843ce4

Please sign in to comment.