Skip to content

Commit e8663dd

Browse files
JU-2094JU-2094
authored
feat: Chat Message History. (#8)
* feat: Chat Message History. * fix: issue with encoding. * chore: refactor. * style: format fix. * chore: refactor user_agent in client. * style: refactor static methods. * style: format fix. * style: fix format.
1 parent 28a9c01 commit e8663dd

File tree

4 files changed

+384
-78
lines changed

4 files changed

+384
-78
lines changed

‎docs/chat_message_history.ipynb

Lines changed: 148 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -1,79 +1,150 @@
11
{
2-
"cells": [
3-
{
4-
"cell_type": "markdown",
5-
"metadata": {},
6-
"source": [
7-
"# Google DATABASE\n",
8-
"\n",
9-
"[Google DATABASE](https://cloud.google.com/DATABASE).\n",
10-
"\n",
11-
"Save chat messages into `DATABASE`."
12-
]
13-
},
14-
{
15-
"cell_type": "markdown",
16-
"metadata": {},
17-
"source": [
18-
"## Pre-reqs"
19-
]
20-
},
21-
{
22-
"cell_type": "code",
23-
"execution_count": null,
24-
"metadata": {
25-
"tags": []
26-
},
27-
"outputs": [],
28-
"source": [
29-
"%pip install PACKAGE_NAME"
30-
]
31-
},
32-
{
33-
"cell_type": "code",
34-
"execution_count": 3,
35-
"metadata": {
36-
"tags": []
37-
},
38-
"outputs": [],
39-
"source": [
40-
"from PACKAGE import LOADER"
41-
]
42-
},
43-
{
44-
"cell_type": "markdown",
45-
"metadata": {},
46-
"source": [
47-
"## Basic Usage"
48-
]
49-
},
50-
{
51-
"cell_type": "code",
52-
"execution_count": null,
53-
"metadata": {},
54-
"outputs": [],
55-
"source": []
56-
}
57-
],
58-
"metadata": {
59-
"kernelspec": {
60-
"display_name": "Python 3 (ipykernel)",
61-
"language": "python",
62-
"name": "python3"
63-
},
64-
"language_info": {
65-
"codemirror_mode": {
66-
"name": "ipython",
67-
"version": 3
68-
},
69-
"file_extension": ".py",
70-
"mimetype": "text/x-python",
71-
"name": "python",
72-
"nbconvert_exporter": "python",
73-
"pygments_lexer": "ipython3",
74-
"version": "3.10.6"
75-
}
76-
},
77-
"nbformat": 4,
78-
"nbformat_minor": 4
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# Cloud Firestore in Datastore Mode\n",
8+
"\n",
9+
"[Cloud Firestore in Datastore Mode](https://cloud.google.com/datastore) is a NoSQL document database build for automatic scaling, high performance and ease of application development.\n",
10+
"\n",
11+
"This package allows you to save chat messages into `Firestore` in Datastore Mode."
12+
]
13+
},
14+
{
15+
"cell_type": "markdown",
16+
"metadata": {},
17+
"source": [
18+
"## Pre-reqs"
19+
]
20+
},
21+
{
22+
"cell_type": "code",
23+
"execution_count": null,
24+
"metadata": {
25+
"tags": []
26+
},
27+
"outputs": [],
28+
"source": [
29+
"%pip install langchain-google-datastore"
30+
]
31+
},
32+
{
33+
"cell_type": "code",
34+
"execution_count": 3,
35+
"metadata": {
36+
"tags": []
37+
},
38+
"outputs": [],
39+
"source": [
40+
"from langchain_google_datastore import DatastoreChatMessageHistory"
41+
]
42+
},
43+
{
44+
"cell_type": "markdown",
45+
"metadata": {},
46+
"source": [
47+
"## Basic Usage"
48+
]
49+
},
50+
{
51+
"cell_type": "code",
52+
"execution_count": null,
53+
"metadata": {},
54+
"outputs": [],
55+
"source": [
56+
"from langchain_google_datastore import DatastoreChatMessageHistory\n",
57+
"\n",
58+
"chat_history = DatastoreChatMessageHistory(\n",
59+
" session_id=\"user-session-id\", kind=\"HistoryMessages\"\n",
60+
")\n",
61+
"\n",
62+
"chat_history.add_user_message(\"Hi!\")\n",
63+
"chat_history.add_ai_message(\"How can I help you?\")"
64+
]
65+
},
66+
{
67+
"cell_type": "code",
68+
"execution_count": null,
69+
"metadata": {},
70+
"outputs": [],
71+
"source": [
72+
"chat_history.messages"
73+
]
74+
},
75+
{
76+
"cell_type": "markdown",
77+
"metadata": {},
78+
"source": [
79+
"## Cleanup\n",
80+
"\n",
81+
"Deleting the history for the session from the database and memory."
82+
]
83+
},
84+
{
85+
"cell_type": "code",
86+
"execution_count": null,
87+
"metadata": {},
88+
"outputs": [],
89+
"source": [
90+
"chat_history.clear()"
91+
]
92+
},
93+
{
94+
"cell_type": "markdown",
95+
"metadata": {},
96+
"source": [
97+
"## Custom Client\n",
98+
"\n",
99+
"The client is created by default using the available environment variables. A [custom client](https://cloud.google.com/python/docs/reference/datastore/latest/client) can be passed to the constructor."
100+
]
101+
},
102+
{
103+
"cell_type": "code",
104+
"execution_count": null,
105+
"metadata": {},
106+
"outputs": [],
107+
"source": [
108+
"from google.auth import compute_engine\n",
109+
"from google.cloud import datastore\n",
110+
"\n",
111+
"client = datastore.Client(\n",
112+
" project=\"project-custom\",\n",
113+
" database=\"non-default-database\",\n",
114+
" credentials=compute_engine.Credentials(),\n",
115+
")\n",
116+
"\n",
117+
"history = DatastoreChatMessageHistory(\n",
118+
" session_id=\"session-id\", kind=\"History\", client=client\n",
119+
")\n",
120+
"\n",
121+
"history.add_user_message(\"New message\")\n",
122+
"\n",
123+
"history.messages\n",
124+
"\n",
125+
"history.clear()"
126+
]
127+
}
128+
],
129+
"metadata": {
130+
"kernelspec": {
131+
"display_name": "Python 3 (ipykernel)",
132+
"language": "python",
133+
"name": "python3"
134+
},
135+
"language_info": {
136+
"codemirror_mode": {
137+
"name": "ipython",
138+
"version": 3
139+
},
140+
"file_extension": ".py",
141+
"mimetype": "text/x-python",
142+
"name": "python",
143+
"nbconvert_exporter": "python",
144+
"pygments_lexer": "ipython3",
145+
"version": "3.10.6"
146+
}
147+
},
148+
"nbformat": 4,
149+
"nbformat_minor": 4
79150
}

‎src/langchain_google_datastore/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from langchain_google_datastore.chat_message_history import DatastoreChatMessageHistory
1516
from langchain_google_datastore.document_loader import DatastoreLoader, DatastoreSaver
1617

17-
__all__ = ["DatastoreLoader", "DatastoreSaver"]
18+
__all__ = ["DatastoreChatMessageHistory", "DatastoreLoader", "DatastoreSaver"]

‎src/langchain_google_datastore/chat_message_history.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import json
18+
from typing import TYPE_CHECKING, Any, Iterator, List, Optional
19+
20+
from google.cloud import datastore
21+
from langchain_core.chat_history import BaseChatMessageHistory
22+
from langchain_core.messages import BaseMessage, messages_from_dict
23+
24+
USER_AGENT = "langchain-google-datastore-python:chat_history"
25+
DEFAULT_KIND = "ChatHistory"
26+
27+
if TYPE_CHECKING:
28+
from google.cloud.datastore import Client # type: ignore
29+
30+
31+
class DatastoreChatMessageHistory(BaseChatMessageHistory):
32+
def __init__(
33+
self,
34+
session_id: str,
35+
kind: str = DEFAULT_KIND,
36+
client: Optional[Client] = None,
37+
) -> None:
38+
"""Chat Message History for Google Cloud Firestore.
39+
Args:
40+
session_id: Arbitrary key that is used to store the messages of a single
41+
chat session. This is the identifier of an entity.
42+
kind: The name of the Datastore kind to write into. This is an optional value
43+
and by default it will use `ChatHistory` as the kind.
44+
client: Client for interacting with the Google Cloud Firestore API.
45+
"""
46+
self.client = client or datastore.Client()
47+
client_agent = self.client._client_info.user_agent
48+
if not client_agent:
49+
self.client._client_info.user_agent = USER_AGENT
50+
elif USER_AGENT not in client_agent:
51+
self.client._client_info.user_agent = " ".join([client_agent, USER_AGENT])
52+
self.session_id = session_id
53+
self.key = self.client.key(kind, session_id)
54+
self.messages: List[BaseMessage] = []
55+
self._load_messages()
56+
57+
def _load_messages(self) -> None:
58+
entity = self.client.get(self.key)
59+
if entity:
60+
data_entity = dict(entity.items())
61+
if "messages" in data_entity:
62+
self.messages = decode_messages(data_entity["messages"])
63+
64+
def add_message(self, message: BaseMessage) -> None:
65+
self.messages.append(message)
66+
self._upsert_messages()
67+
68+
def _upsert_messages(self) -> None:
69+
entity = self.client.entity(self.key)
70+
entity["messages"] = encode_messages(self.messages)
71+
self.client.put(entity)
72+
73+
def clear(self) -> None:
74+
self.messages = []
75+
self.client.delete(self.key)
76+
77+
78+
def encode_messages(messages: List[BaseMessage]) -> List[bytes]:
79+
return [str.encode(m.json()) for m in messages]
80+
81+
82+
def decode_messages(messages: List[bytes]) -> List[BaseMessage]:
83+
dict_messages = [json.loads(m.decode()) for m in messages]
84+
return messages_from_dict([{"type": m["type"], "data": m} for m in dict_messages])

0 commit comments

Comments
 (0)