github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-18 02:33:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

added cosmos kwargs option (#5292)

Browse Source 
# Added the ability to pass kwargs to cosmos client constructor

The cosmos client has a ton of options that can be set, so allowing
those to be passed to the constructor from the chat memory constructor
with this PR.
This commit is contained in:
Eduard van Valkenburg 2023-05-28 06:19:40 +02:00 committed by GitHub
parent 881dfe8179
commit 1daa7068b2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
langchain/memory/chat_message_histories/cosmos_db.py
View File

@ -33,6 +33,7 @@ class CosmosDBChatMessageHistory(BaseChatMessageHistory):
credential: Any = None, credential: Any = None,
connection_string: Optional[str] = None, connection_string: Optional[str] = None,
ttl: Optional[int] = None, ttl: Optional[int] = None,
cosmos_client_kwargs: Optional[dict] = None,
): ):
""" """
Initializes a new instance of the CosmosDBChatMessageHistory class. Initializes a new instance of the CosmosDBChatMessageHistory class.
@ -50,6 +51,7 @@ class CosmosDBChatMessageHistory(BaseChatMessageHistory):
:param credential: The credential to use to authenticate to Azure Cosmos DB. :param credential: The credential to use to authenticate to Azure Cosmos DB.
:param connection_string: The connection string to use to authenticate. :param connection_string: The connection string to use to authenticate.
:param ttl: The time to live (in seconds) to use for documents in the container. :param ttl: The time to live (in seconds) to use for documents in the container.
:param cosmos_client_kwargs: Additional kwargs to pass to the CosmosClient.
""" """
self.cosmos_endpoint = cosmos_endpoint self.cosmos_endpoint = cosmos_endpoint
self.cosmos_database = cosmos_database self.cosmos_database = cosmos_database
@ -71,11 +73,14 @@ class CosmosDBChatMessageHistory(BaseChatMessageHistory):
) from exc ) from exc
if self.credential: if self.credential:
self._client = CosmosClient( self._client = CosmosClient(
url=self.cosmos_endpoint, credential=self.credential url=self.cosmos_endpoint,
credential=self.credential,
**cosmos_client_kwargs or {},
) )
elif self.conn_string: elif self.conn_string:
self._client = CosmosClient.from_connection_string( self._client = CosmosClient.from_connection_string(
conn_str=self.conn_string conn_str=self.conn_string,
**cosmos_client_kwargs or {},
) )
else: else:
raise ValueError("Either a connection string or a credential must be set.") raise ValueError("Either a connection string or a credential must be set.")