Skip to content

Commit 1843ce4

Browse files
amotlamotl
committed
CrateDB memory: Add conversational memory support
The implementation is based on the generic `SQLChatMessageHistory`.
1 parent 571ce1f commit 1843ce4

File tree

6 files changed

+679
-2
lines changed

6 files changed

+679
-2
lines changed

‎docs/docs_skeleton/vercel.json

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2632,6 +2632,14 @@
26322632
"source": "/docs/modules/memory/integrations/cassandra_chat_message_history",
26332633
"destination": "/docs/integrations/memory/cassandra_chat_message_history"
26342634
},
2635+
{
2636+
"source": "/en/latest/modules/memory/examples/cratedb_chat_message_history.html",
2637+
"destination": "/docs/integrations/memory/cratedb_chat_message_history"
2638+
},
2639+
{
2640+
"source": "/docs/modules/memory/integrations/cratedb_chat_message_history",
2641+
"destination": "/docs/integrations/memory/cratedb_chat_message_history"
2642+
},
26352643
{
26362644
"source": "/en/latest/modules/memory/examples/dynamodb_chat_message_history.html",
26372645
"destination": "/docs/integrations/memory/dynamodb_chat_message_history"

‎docs/extras/integrations/memory/cratedb_chat_message_history.ipynb

Lines changed: 357 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,357 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"source": [
6+
"# CrateDB Chat Message History\n",
7+
"\n",
8+
"This notebook demonstrates how to use the `CrateDBChatMessageHistory`\n",
9+
"to manage chat history in CrateDB, for supporting conversational memory."
10+
],
11+
"metadata": {
12+
"collapsed": false
13+
},
14+
"id": "f22eab3f84cbeb37"
15+
},
16+
{
17+
"cell_type": "markdown",
18+
"source": [
19+
"## Prerequisites"
20+
],
21+
"metadata": {
22+
"collapsed": false
23+
}
24+
},
25+
{
26+
"cell_type": "code",
27+
"execution_count": null,
28+
"outputs": [],
29+
"source": [
30+
"!#pip install 'crate[sqlalchemy]' 'langchain'"
31+
],
32+
"metadata": {
33+
"collapsed": false
34+
}
35+
},
36+
{
37+
"cell_type": "markdown",
38+
"source": [
39+
"## Configuration\n",
40+
"\n",
41+
"To use the storage wrapper, you will need to configure two details.\n",
42+
"\n",
43+
"1. Session Id - a unique identifier of the session, like user name, email, chat id etc.\n",
44+
"2. Database connection string: An SQLAlchemy-compatible URI that specifies the database\n",
45+
" connection. It will be passed to SQLAlchemy create_engine function."
46+
],
47+
"metadata": {
48+
"collapsed": false
49+
},
50+
"id": "f8f2830ee9ca1e01"
51+
},
52+
{
53+
"cell_type": "code",
54+
"execution_count": 52,
55+
"outputs": [],
56+
"source": [
57+
"from langchain.memory.chat_message_histories import CrateDBChatMessageHistory\n",
58+
"\n",
59+
"CONNECTION_STRING = \"crate://crate@localhost:4200/?schema=example\"\n",
60+
"\n",
61+
"chat_message_history = CrateDBChatMessageHistory(\n",
62+
"\tsession_id=\"test_session\",\n",
63+
"\tconnection_string=CONNECTION_STRING\n",
64+
")"
65+
],
66+
"metadata": {
67+
"collapsed": false
68+
}
69+
},
70+
{
71+
"cell_type": "markdown",
72+
"source": [
73+
"## Basic Usage"
74+
],
75+
"metadata": {
76+
"collapsed": false
77+
}
78+
},
79+
{
80+
"cell_type": "code",
81+
"execution_count": 53,
82+
"outputs": [],
83+
"source": [
84+
"chat_message_history.add_user_message(\"Hello\")\n",
85+
"chat_message_history.add_ai_message(\"Hi\")"
86+
],
87+
"metadata": {
88+
"collapsed": false,
89+
"ExecuteTime": {
90+
"end_time": "2023-08-28T10:04:38.077748Z",
91+
"start_time": "2023-08-28T10:04:36.105894Z"
92+
}
93+
},
94+
"id": "4576e914a866fb40"
95+
},
96+
{
97+
"cell_type": "code",
98+
"execution_count": 61,
99+
"outputs": [
100+
{
101+
"data": {
102+
"text/plain": "[HumanMessage(content='Hello', additional_kwargs={}, example=False),\n AIMessage(content='Hi', additional_kwargs={}, example=False)]"
103+
},
104+
"execution_count": 61,
105+
"metadata": {},
106+
"output_type": "execute_result"
107+
}
108+
],
109+
"source": [
110+
"chat_message_history.messages"
111+
],
112+
"metadata": {
113+
"collapsed": false,
114+
"ExecuteTime": {
115+
"end_time": "2023-08-28T10:04:38.929396Z",
116+
"start_time": "2023-08-28T10:04:38.915727Z"
117+
}
118+
},
119+
"id": "b476688cbb32ba90"
120+
},
121+
{
122+
"cell_type": "markdown",
123+
"source": [
124+
"## Custom Storage Model\n",
125+
"\n",
126+
"The default data model, which stores information about conversation messages only\n",
127+
"has two slots for storing message details, the session id, and the message dictionary.\n",
128+
"\n",
129+
"If you want to store additional information, like message date, author, language etc.,\n",
130+
"please provide an implementation for a custom message converter.\n",
131+
"\n",
132+
"This example demonstrates how to create a custom message converter, by implementing\n",
133+
"the `BaseMessageConverter` interface."
134+
],
135+
"metadata": {
136+
"collapsed": false
137+
},
138+
"id": "2e5337719d5614fd"
139+
},
140+
{
141+
"cell_type": "code",
142+
"execution_count": 55,
143+
"outputs": [],
144+
"source": [
145+
"from datetime import datetime\n",
146+
"from typing import Any\n",
147+
"\n",
148+
"from langchain.memory.chat_message_histories.cratedb import generate_autoincrement_identifier\n",
149+
"from langchain.memory.chat_message_histories.sql import BaseMessageConverter\n",
150+
"from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage\n",
151+
"\n",
152+
"import sqlalchemy as sa\n",
153+
"from sqlalchemy.orm import declarative_base\n",
154+
"\n",
155+
"\n",
156+
"Base = declarative_base()\n",
157+
"\n",
158+
"\n",
159+
"class CustomMessage(Base):\n",
160+
"\t__tablename__ = \"custom_message_store\"\n",
161+
"\n",
162+
"\tid = sa.Column(sa.BigInteger, primary_key=True, default=generate_autoincrement_identifier)\n",
163+
"\tsession_id = sa.Column(sa.Text)\n",
164+
"\ttype = sa.Column(sa.Text)\n",
165+
"\tcontent = sa.Column(sa.Text)\n",
166+
"\tcreated_at = sa.Column(sa.DateTime)\n",
167+
"\tauthor_email = sa.Column(sa.Text)\n",
168+
"\n",
169+
"\n",
170+
"class CustomMessageConverter(BaseMessageConverter):\n",
171+
"\tdef __init__(self, author_email: str):\n",
172+
"\t\tself.author_email = author_email\n",
173+
"\t\n",
174+
"\tdef from_sql_model(self, sql_message: Any) -> BaseMessage:\n",
175+
"\t\tif sql_message.type == \"human\":\n",
176+
"\t\t\treturn HumanMessage(\n",
177+
"\t\t\t\tcontent=sql_message.content,\n",
178+
"\t\t\t)\n",
179+
"\t\telif sql_message.type == \"ai\":\n",
180+
"\t\t\treturn AIMessage(\n",
181+
"\t\t\t\tcontent=sql_message.content,\n",
182+
"\t\t\t)\n",
183+
"\t\telif sql_message.type == \"system\":\n",
184+
"\t\t\treturn SystemMessage(\n",
185+
"\t\t\t\tcontent=sql_message.content,\n",
186+
"\t\t\t)\n",
187+
"\t\telse:\n",
188+
"\t\t\traise ValueError(f\"Unknown message type: {sql_message.type}\")\n",
189+
"\t\n",
190+
"\tdef to_sql_model(self, message: BaseMessage, session_id: str) -> Any:\n",
191+
"\t\tnow = datetime.now()\n",
192+
"\t\treturn CustomMessage(\n",
193+
"\t\t\tsession_id=session_id,\n",
194+
"\t\t\ttype=message.type,\n",
195+
"\t\t\tcontent=message.content,\n",
196+
"\t\t\tcreated_at=now,\n",
197+
"\t\t\tauthor_email=self.author_email\n",
198+
"\t\t)\n",
199+
"\t\n",
200+
"\tdef get_sql_model_class(self) -> Any:\n",
201+
"\t\treturn CustomMessage\n",
202+
"\n",
203+
"\n",
204+
"if __name__ == \"__main__\":\n",
205+
"\n",
206+
"\tBase.metadata.drop_all(bind=sa.create_engine(CONNECTION_STRING))\n",
207+
"\n",
208+
"\tchat_message_history = CrateDBChatMessageHistory(\n",
209+
"\t\tsession_id=\"test_session\",\n",
210+
"\t\tconnection_string=CONNECTION_STRING,\n",
211+
"\t\tcustom_message_converter=CustomMessageConverter(\n",
212+
"\t\t\tauthor_email=\"test@example.com\"\n",
213+
"\t\t)\n",
214+
"\t)\n",
215+
"\n",
216+
"\tchat_message_history.add_user_message(\"Hello\")\n",
217+
"\tchat_message_history.add_ai_message(\"Hi\")"
218+
],
219+
"metadata": {
220+
"collapsed": false,
221+
"ExecuteTime": {
222+
"end_time": "2023-08-28T10:04:41.510498Z",
223+
"start_time": "2023-08-28T10:04:41.494912Z"
224+
}
225+
},
226+
"id": "fdfde84c07d071bb"
227+
},
228+
{
229+
"cell_type": "code",
230+
"execution_count": 60,
231+
"outputs": [
232+
{
233+
"data": {
234+
"text/plain": "[HumanMessage(content='Hello', additional_kwargs={}, example=False),\n AIMessage(content='Hi', additional_kwargs={}, example=False)]"
235+
},
236+
"execution_count": 60,
237+
"metadata": {},
238+
"output_type": "execute_result"
239+
}
240+
],
241+
"source": [
242+
"chat_message_history.messages"
243+
],
244+
"metadata": {
245+
"collapsed": false,
246+
"ExecuteTime": {
247+
"end_time": "2023-08-28T10:04:43.497990Z",
248+
"start_time": "2023-08-28T10:04:43.492517Z"
249+
}
250+
},
251+
"id": "4a6a54d8a9e2856f"
252+
},
253+
{
254+
"cell_type": "markdown",
255+
"source": [
256+
"## Custom Name for Session Column\n",
257+
"\n",
258+
"The session id, a unique token identifying the session, is an important property of\n",
259+
"this subsystem. If your database table stores it in a different column, you can use\n",
260+
"the `session_id_field_name` keyword argument to adjust the name correspondingly."
261+
],
262+
"metadata": {
263+
"collapsed": false
264+
},
265+
"id": "622aded629a1adeb"
266+
},
267+
{
268+
"cell_type": "code",
269+
"execution_count": 57,
270+
"outputs": [],
271+
"source": [
272+
"import json\n",
273+
"import typing as t\n",
274+
"\n",
275+
"from langchain.memory.chat_message_histories.cratedb import generate_autoincrement_identifier, CrateDBMessageConverter\n",
276+
"from langchain.schema import _message_to_dict\n",
277+
"\n",
278+
"\n",
279+
"Base = declarative_base()\n",
280+
"\n",
281+
"class MessageWithDifferentSessionIdColumn(Base):\n",
282+
"\t__tablename__ = \"message_store_different_session_id\"\n",
283+
"\tid = sa.Column(sa.BigInteger, primary_key=True, default=generate_autoincrement_identifier)\n",
284+
"\tcustom_session_id = sa.Column(sa.Text)\n",
285+
"\tmessage = sa.Column(sa.Text)\n",
286+
"\n",
287+
"\n",
288+
"class CustomMessageConverterWithDifferentSessionIdColumn(CrateDBMessageConverter):\n",
289+
" def __init__(self):\n",
290+
" self.model_class = MessageWithDifferentSessionIdColumn\n",
291+
"\n",
292+
" def to_sql_model(self, message: BaseMessage, custom_session_id: str) -> t.Any:\n",
293+
" return self.model_class(\n",
294+
" custom_session_id=custom_session_id, message=json.dumps(_message_to_dict(message))\n",
295+
" )\n",
296+
"\n",
297+
"\n",
298+
"if __name__ == \"__main__\":\n",
299+
"\tBase.metadata.drop_all(bind=sa.create_engine(CONNECTION_STRING))\n",
300+
"\n",
301+
"\tchat_message_history = CrateDBChatMessageHistory(\n",
302+
"\t\tsession_id=\"test_session\",\n",
303+
"\t\tconnection_string=CONNECTION_STRING,\n",
304+
"\t\tcustom_message_converter=CustomMessageConverterWithDifferentSessionIdColumn(),\n",
305+
"\t\tsession_id_field_name=\"custom_session_id\",\n",
306+
"\t)\n",
307+
"\n",
308+
"\tchat_message_history.add_user_message(\"Hello\")\n",
309+
"\tchat_message_history.add_ai_message(\"Hi\")"
310+
],
311+
"metadata": {
312+
"collapsed": false
313+
}
314+
},
315+
{
316+
"cell_type": "code",
317+
"execution_count": 58,
318+
"outputs": [
319+
{
320+
"data": {
321+
"text/plain": "[HumanMessage(content='Hello', additional_kwargs={}, example=False),\n AIMessage(content='Hi', additional_kwargs={}, example=False)]"
322+
},
323+
"execution_count": 58,
324+
"metadata": {},
325+
"output_type": "execute_result"
326+
}
327+
],
328+
"source": [
329+
"chat_message_history.messages"
330+
],
331+
"metadata": {
332+
"collapsed": false
333+
}
334+
}
335+
],
336+
"metadata": {
337+
"kernelspec": {
338+
"display_name": "Python 3",
339+
"language": "python",
340+
"name": "python3"
341+
},
342+
"language_info": {
343+
"codemirror_mode": {
344+
"name": "ipython",
345+
"version": 2
346+
},
347+
"file_extension": ".py",
348+
"mimetype": "text/x-python",
349+
"name": "python",
350+
"nbconvert_exporter": "python",
351+
"pygments_lexer": "ipython2",
352+
"version": "2.7.6"
353+
}
354+
},
355+
"nbformat": 4,
356+
"nbformat_minor": 5
357+
}

0 commit comments

Comments
 (0)