Skip to content

Commit

Permalink
docs: Fix AstraDBChatMessageHistory docstrings (#17740)
Browse files Browse the repository at this point in the history
  • Loading branch information
cbornet committed Feb 19, 2024
1 parent 86ae48b commit 6275d8b
Showing 1 changed file with 6 additions and 12 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -26,16 +26,16 @@
class AstraDBChatMessageHistory(BaseChatMessageHistory):
"""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
of a single chat session.
collection_name (str): name of the Astra DB collection to create/use.
token (Optional[str]): API token for Astra DB usage.
api_endpoint (Optional[str]): full URL to the API endpoint,
collection_name: name of the Astra DB collection to create/use.
token: API token for Astra DB usage.
api_endpoint: full URL to the API endpoint,
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.
namespace (Optional[str]): namespace (aka keyspace) where the
namespace: namespace (aka keyspace) where the
collection is created. Defaults to the database's "default namespace".
"""

Expand All @@ -51,7 +51,6 @@ def __init__(
setup_mode: SetupMode = SetupMode.SYNC,
pre_delete_collection: bool = False,
) -> None:
"""Create an Astra DB chat message history."""
self.astra_env = _AstraDBCollectionEnvironment(
collection_name=collection_name,
token=token,
Expand Down Expand Up @@ -96,7 +95,6 @@ def messages(self, messages: List[BaseMessage]) -> None:
raise NotImplementedError("Use add_messages instead")

async def aget_messages(self) -> List[BaseMessage]:
"""Retrieve all session messages from DB"""
await self.astra_env.aensure_db_setup()
docs = self.async_collection.paginated_find(
filter={
Expand All @@ -117,7 +115,6 @@ async def aget_messages(self) -> List[BaseMessage]:
return messages

def add_messages(self, messages: Sequence[BaseMessage]) -> None:
"""Write a message to the table"""
self.astra_env.ensure_db_setup()
docs = [
{
Expand All @@ -130,7 +127,6 @@ def add_messages(self, messages: Sequence[BaseMessage]) -> None:
self.collection.chunked_insert_many(docs)

async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:
"""Write a message to the table"""
await self.astra_env.aensure_db_setup()
docs = [
{
Expand All @@ -143,11 +139,9 @@ async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:
await self.async_collection.chunked_insert_many(docs)

def clear(self) -> None:
"""Clear session memory from DB"""
self.astra_env.ensure_db_setup()
self.collection.delete_many(filter={"session_id": self.session_id})

async def aclear(self) -> None:
"""Clear session memory from DB"""
await self.astra_env.aensure_db_setup()
await self.async_collection.delete_many(filter={"session_id": self.session_id})

0 comments on commit 6275d8b

Please sign in to comment.