docs: Fix AstraDBChatMessageHistory docstrings (#17740)

This commit is contained in:
Christophe Bornet 2024-02-19 19:47:38 +01:00 committed by GitHub
parent 86ae48b781
commit 6275d8b1bf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -26,16 +26,16 @@ DEFAULT_COLLECTION_NAME = "langchain_message_store"
class AstraDBChatMessageHistory(BaseChatMessageHistory): class AstraDBChatMessageHistory(BaseChatMessageHistory):
"""Chat message history that stores history in Astra DB. """Chat message history that stores history in Astra DB.
Args (only keyword-arguments accepted): Args:
session_id: arbitrary key that is used to store the messages session_id: arbitrary key that is used to store the messages
of a single chat session. of a single chat session.
collection_name (str): name of the Astra DB collection to create/use. collection_name: name of the Astra DB collection to create/use.
token (Optional[str]): API token for Astra DB usage. token: API token for Astra DB usage.
api_endpoint (Optional[str]): full URL to the API endpoint, api_endpoint: full URL to the API endpoint,
such as "https://<DB-ID>-us-east1.apps.astra.datastax.com". such as "https://<DB-ID>-us-east1.apps.astra.datastax.com".
astra_db_client (Optional[Any]): *alternative to token+api_endpoint*, astra_db_client: *alternative to token+api_endpoint*,
you can pass an already-created 'astrapy.db.AstraDB' instance. you can pass an already-created 'astrapy.db.AstraDB' instance.
namespace (Optional[str]): namespace (aka keyspace) where the namespace: namespace (aka keyspace) where the
collection is created. Defaults to the database's "default namespace". collection is created. Defaults to the database's "default namespace".
""" """
@ -51,7 +51,6 @@ class AstraDBChatMessageHistory(BaseChatMessageHistory):
setup_mode: SetupMode = SetupMode.SYNC, setup_mode: SetupMode = SetupMode.SYNC,
pre_delete_collection: bool = False, pre_delete_collection: bool = False,
) -> None: ) -> None:
"""Create an Astra DB chat message history."""
self.astra_env = _AstraDBCollectionEnvironment( self.astra_env = _AstraDBCollectionEnvironment(
collection_name=collection_name, collection_name=collection_name,
token=token, token=token,
@ -96,7 +95,6 @@ class AstraDBChatMessageHistory(BaseChatMessageHistory):
raise NotImplementedError("Use add_messages instead") raise NotImplementedError("Use add_messages instead")
async def aget_messages(self) -> List[BaseMessage]: async def aget_messages(self) -> List[BaseMessage]:
"""Retrieve all session messages from DB"""
await self.astra_env.aensure_db_setup() await self.astra_env.aensure_db_setup()
docs = self.async_collection.paginated_find( docs = self.async_collection.paginated_find(
filter={ filter={
@ -117,7 +115,6 @@ class AstraDBChatMessageHistory(BaseChatMessageHistory):
return messages return messages
def add_messages(self, messages: Sequence[BaseMessage]) -> None: def add_messages(self, messages: Sequence[BaseMessage]) -> None:
"""Write a message to the table"""
self.astra_env.ensure_db_setup() self.astra_env.ensure_db_setup()
docs = [ docs = [
{ {
@ -130,7 +127,6 @@ class AstraDBChatMessageHistory(BaseChatMessageHistory):
self.collection.chunked_insert_many(docs) self.collection.chunked_insert_many(docs)
async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None: async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:
"""Write a message to the table"""
await self.astra_env.aensure_db_setup() await self.astra_env.aensure_db_setup()
docs = [ docs = [
{ {
@ -143,11 +139,9 @@ class AstraDBChatMessageHistory(BaseChatMessageHistory):
await self.async_collection.chunked_insert_many(docs) await self.async_collection.chunked_insert_many(docs)
def clear(self) -> None: def clear(self) -> None:
"""Clear session memory from DB"""
self.astra_env.ensure_db_setup() self.astra_env.ensure_db_setup()
self.collection.delete_many(filter={"session_id": self.session_id}) self.collection.delete_many(filter={"session_id": self.session_id})
async def aclear(self) -> None: async def aclear(self) -> None:
"""Clear session memory from DB"""
await self.astra_env.aensure_db_setup() await self.astra_env.aensure_db_setup()
await self.async_collection.delete_many(filter={"session_id": self.session_id}) await self.async_collection.delete_many(filter={"session_id": self.session_id})