Skip to content

Commit 7787d7d

Browse files
dishaprakashdishaprakash
authored
feat: Add Postgres Chat Store (#40)
* feat: Add Postgres Chat Store * Linter fix
1 parent 2b14f5a commit 7787d7d

File tree

3 files changed

+673
-0
lines changed

3 files changed

+673
-0
lines changed

‎src/llama_index_cloud_sql_pg/__init__.py

Lines changed: 2 additions & 0 deletions
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 .chat_store import PostgresChatStore
1516
from .document_store import PostgresDocumentStore
1617
from .engine import Column, PostgresEngine
1718
from .index_store import PostgresIndexStore
@@ -20,6 +21,7 @@
2021

2122
_all = [
2223
"Column",
24+
"PostgresChatStore",
2325
"PostgresEngine",
2426
"PostgresDocumentStore",
2527
"PostgresIndexStore",

‎src/llama_index_cloud_sql_pg/chat_store.py

Lines changed: 289 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,289 @@
1+
# Copyright 2025 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+
from typing import List, Optional
18+
19+
from llama_index.core.llms import ChatMessage
20+
from llama_index.core.storage.chat_store.base import BaseChatStore
21+
22+
from .async_chat_store import AsyncPostgresChatStore
23+
from .engine import PostgresEngine
24+
25+
26+
class PostgresChatStore(BaseChatStore):
27+
"""Chat Store Table stored in an Cloud SQL for PostgreSQL database."""
28+
29+
__create_key = object()
30+
31+
def __init__(
32+
self, key: object, engine: PostgresEngine, chat_store: AsyncPostgresChatStore
33+
):
34+
"""PostgresChatStore constructor.
35+
36+
Args:
37+
key (object): Key to prevent direct constructor usage.
38+
engine (PostgresEngine): Database connection pool.
39+
chat_store (AsyncPostgresChatStore): The async only IndexStore implementation
40+
41+
Raises:
42+
Exception: If constructor is directly called by the user.
43+
"""
44+
if key != PostgresChatStore.__create_key:
45+
raise Exception(
46+
"Only create class through 'create' or 'create_sync' methods!"
47+
)
48+
49+
# Delegate to Pydantic's __init__
50+
super().__init__()
51+
self._engine = engine
52+
self.__chat_store = chat_store
53+
54+
@classmethod
55+
async def create(
56+
cls,
57+
engine: PostgresEngine,
58+
table_name: str,
59+
schema_name: str = "public",
60+
) -> PostgresChatStore:
61+
"""Create a new PostgresChatStore instance.
62+
63+
Args:
64+
engine (PostgresEngine): Postgres engine to use.
65+
table_name (str): Table name that stores the chat store.
66+
schema_name (str): The schema name where the table is located. Defaults to "public"
67+
68+
Raises:
69+
ValueError: If the table provided does not contain required schema.
70+
71+
Returns:
72+
PostgresChatStore: A newly created instance of PostgresChatStore.
73+
"""
74+
coro = AsyncPostgresChatStore.create(engine, table_name, schema_name)
75+
chat_store = await engine._run_as_async(coro)
76+
return cls(cls.__create_key, engine, chat_store)
77+
78+
@classmethod
79+
def create_sync(
80+
cls,
81+
engine: PostgresEngine,
82+
table_name: str,
83+
schema_name: str = "public",
84+
) -> PostgresChatStore:
85+
"""Create a new PostgresChatStore sync instance.
86+
87+
Args:
88+
engine (PostgresEngine): Postgres engine to use.
89+
table_name (str): Table name that stores the chat store.
90+
schema_name (str): The schema name where the table is located. Defaults to "public"
91+
92+
Raises:
93+
ValueError: If the table provided does not contain required schema.
94+
95+
Returns:
96+
PostgresChatStore: A newly created instance of PostgresChatStore.
97+
"""
98+
coro = AsyncPostgresChatStore.create(engine, table_name, schema_name)
99+
chat_store = engine._run_as_sync(coro)
100+
return cls(cls.__create_key, engine, chat_store)
101+
102+
@classmethod
103+
def class_name(cls) -> str:
104+
"""Get class name."""
105+
return "PostgresChatStore"
106+
107+
async def aset_messages(self, key: str, messages: List[ChatMessage]) -> None:
108+
"""Asynchronously sets the chat messages for a specific key.
109+
110+
Args:
111+
key (str): A unique identifier for the chat.
112+
messages (List[ChatMessage]): A list of `ChatMessage` objects to upsert.
113+
114+
Returns:
115+
None
116+
117+
"""
118+
return await self._engine._run_as_async(
119+
self.__chat_store.aset_messages(key=key, messages=messages)
120+
)
121+
122+
async def aget_messages(self, key: str) -> List[ChatMessage]:
123+
"""Asynchronously retrieves the chat messages associated with a specific key.
124+
125+
Args:
126+
key (str): A unique identifier for which the messages are to be retrieved.
127+
128+
Returns:
129+
List[ChatMessage]: A list of `ChatMessage` objects associated with the provided key.
130+
If no messages are found, an empty list is returned.
131+
"""
132+
return await self._engine._run_as_async(
133+
self.__chat_store.aget_messages(key=key)
134+
)
135+
136+
async def async_add_message(self, key: str, message: ChatMessage) -> None:
137+
"""Asynchronously adds a new chat message to the specified key.
138+
139+
Args:
140+
key (str): A unique identifierfor the chat to which the message is added.
141+
message (ChatMessage): The `ChatMessage` object that is to be added.
142+
143+
Returns:
144+
None
145+
"""
146+
return await self._engine._run_as_async(
147+
self.__chat_store.async_add_message(key=key, message=message)
148+
)
149+
150+
async def adelete_messages(self, key: str) -> Optional[List[ChatMessage]]:
151+
"""Asynchronously deletes the chat messages associated with a specific key.
152+
153+
Args:
154+
key (str): A unique identifier for the chat whose messages are to be deleted.
155+
156+
Returns:
157+
Optional[List[ChatMessage]]: A list of `ChatMessage` objects that were deleted, or `None` if no messages
158+
were associated with the key or could be deleted.
159+
"""
160+
return await self._engine._run_as_async(
161+
self.__chat_store.adelete_messages(key=key)
162+
)
163+
164+
async def adelete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
165+
"""Asynchronously deletes a specific chat message by index from the messages associated with a given key.
166+
167+
Args:
168+
key (str): A unique identifier for the chat whose messages are to be deleted.
169+
idx (int): The index of the `ChatMessage` to be deleted from the list of messages.
170+
171+
Returns:
172+
Optional[ChatMessage]: The `ChatMessage` object that was deleted, or `None` if no message
173+
was associated with the key or could be deleted.
174+
"""
175+
return await self._engine._run_as_async(
176+
self.__chat_store.adelete_message(key=key, idx=idx)
177+
)
178+
179+
async def adelete_last_message(self, key: str) -> Optional[ChatMessage]:
180+
"""Asynchronously deletes the last chat message associated with a given key.
181+
182+
Args:
183+
key (str): A unique identifier for the chat whose message is to be deleted.
184+
185+
Returns:
186+
Optional[ChatMessage]: The `ChatMessage` object that was deleted, or `None` if no message
187+
was associated with the key or could be deleted.
188+
"""
189+
return await self._engine._run_as_async(
190+
self.__chat_store.adelete_last_message(key=key)
191+
)
192+
193+
async def aget_keys(self) -> List[str]:
194+
"""Asynchronously retrieves a list of all keys.
195+
196+
Returns:
197+
Optional[str]: A list of strings representing the keys. If no keys are found, an empty list is returned.
198+
"""
199+
return await self._engine._run_as_async(self.__chat_store.aget_keys())
200+
201+
def set_messages(self, key: str, messages: List[ChatMessage]) -> None:
202+
"""Synchronously sets the chat messages for a specific key.
203+
204+
Args:
205+
key (str): A unique identifier for the chat.
206+
messages (List[ChatMessage]): A list of `ChatMessage` objects to upsert.
207+
208+
Returns:
209+
None
210+
211+
"""
212+
return self._engine._run_as_sync(
213+
self.__chat_store.aset_messages(key=key, messages=messages)
214+
)
215+
216+
def get_messages(self, key: str) -> List[ChatMessage]:
217+
"""Synchronously retrieves the chat messages associated with a specific key.
218+
219+
Args:
220+
key (str): A unique identifier for which the messages are to be retrieved.
221+
222+
Returns:
223+
List[ChatMessage]: A list of `ChatMessage` objects associated with the provided key.
224+
If no messages are found, an empty list is returned.
225+
"""
226+
return self._engine._run_as_sync(self.__chat_store.aget_messages(key=key))
227+
228+
def add_message(self, key: str, message: ChatMessage) -> None:
229+
"""Synchronously adds a new chat message to the specified key.
230+
231+
Args:
232+
key (str): A unique identifierfor the chat to which the message is added.
233+
message (ChatMessage): The `ChatMessage` object that is to be added.
234+
235+
Returns:
236+
None
237+
"""
238+
return self._engine._run_as_sync(
239+
self.__chat_store.async_add_message(key=key, message=message)
240+
)
241+
242+
def delete_messages(self, key: str) -> Optional[List[ChatMessage]]:
243+
"""Synchronously deletes the chat messages associated with a specific key.
244+
245+
Args:
246+
key (str): A unique identifier for the chat whose messages are to be deleted.
247+
248+
Returns:
249+
Optional[List[ChatMessage]]: A list of `ChatMessage` objects that were deleted, or `None` if no messages
250+
were associated with the key or could be deleted.
251+
"""
252+
return self._engine._run_as_sync(self.__chat_store.adelete_messages(key=key))
253+
254+
def delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
255+
"""Synchronously deletes a specific chat message by index from the messages associated with a given key.
256+
257+
Args:
258+
key (str): A unique identifier for the chat whose messages are to be deleted.
259+
idx (int): The index of the `ChatMessage` to be deleted from the list of messages.
260+
261+
Returns:
262+
Optional[ChatMessage]: The `ChatMessage` object that was deleted, or `None` if no message
263+
was associated with the key or could be deleted.
264+
"""
265+
return self._engine._run_as_sync(
266+
self.__chat_store.adelete_message(key=key, idx=idx)
267+
)
268+
269+
def delete_last_message(self, key: str) -> Optional[ChatMessage]:
270+
"""Synchronously deletes the last chat message associated with a given key.
271+
272+
Args:
273+
key (str): A unique identifier for the chat whose message is to be deleted.
274+
275+
Returns:
276+
Optional[ChatMessage]: The `ChatMessage` object that was deleted, or `None` if no message
277+
was associated with the key or could be deleted.
278+
"""
279+
return self._engine._run_as_sync(
280+
self.__chat_store.adelete_last_message(key=key)
281+
)
282+
283+
def get_keys(self) -> List[str]:
284+
"""Synchronously retrieves a list of all keys.
285+
286+
Returns:
287+
Optional[str]: A list of strings representing the keys. If no keys are found, an empty list is returned.
288+
"""
289+
return self._engine._run_as_sync(self.__chat_store.aget_keys())

0 commit comments

Comments
 (0)